Update app.py

app.py

@@ -2,39 +2,59 @@ import gradio as gr
 import cv2
 import numpy as np
 import torch
-from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForObjectDetection, AutoProcessor
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    AutoModelForObjectDetection,
+    AutoProcessor,
+    pipeline
+)
 from PIL import Image, ImageDraw
 import matplotlib.pyplot as plt
 import pandas as pd
 import warnings
-from ultralytics import YOLO  # Import YOLO from the ultralytics package
+import io
+import base64
 
 warnings.filterwarnings("ignore")
 
 # Constants
-MODEL_NAME = "google/flan-t5-large"
-YOLO_MODEL_PATH = "yolov8n.pt"  # Load the pre-trained YOLOv8 model (you can change this to your desired model)
 MAX_WIDTH = 800
 MAX_HEIGHT = 600
 
-# Load models
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
-yolo_model = YOLO(YOLO_MODEL_PATH)  # Load YOLOv8 model
+# Load models (Update these with your actual model paths/names)
+LLM_MODEL_NAME = "meta-llama/Llama-3-70B-Instruct"
+DETECTION_MODEL = "facebook/detr-resnet-50"
+
+# Initialize device
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# System Prompt Template
-SYSTEM_PROMPT = """
+# Load detection model and processor
+detection_processor = AutoProcessor.from_pretrained(DETECTION_MODEL)
+detection_model = AutoModelForObjectDetection.from_pretrained(DETECTION_MODEL).to(device)
+
+# Load LLM components
+llm_tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_NAME)
+llm_model = AutoModelForCausalLM.from_pretrained(
+    LLM_MODEL_NAME,
+    torch_dtype=torch.bfloat16,
+    device_map="auto"
+)
+
+# System Prompt Template for LLAMA
+SYSTEM_PROMPT = """<|begin_of_text|><|start_header_id|>system<|end_header_id|>
 You are a professional financial analyst specializing in cryptocurrency technical analysis.
-Analyze the following chart elements and provide a detailed report based on the user's question:
+Analyze the following chart elements and provide a detailed report:
+<|eot_id|><|start_header_id|>user<|end_header_id|>
+Chart Elements Detected:
+- Support/Resistance: {support_resistance}
+- Trendlines: {trendlines}
+- Patterns: {patterns}
+- Candlestick formations: {candlesticks}
 
-### Chart Elements:
-Support/Resistance: {support_resistance}
-Trendlines: {trendlines}
-Patterns: {patterns}
-Candlestick formations: {candlesticks}
+User Question: {question}
 
-### User Question:
-{question}
+Provide analysis covering:
 1. Trend analysis (primary and secondary trends)
 2. Key support/resistance levels
 3. Detected chart patterns
@@ -43,16 +63,21 @@ Candlestick formations: {candlesticks}
 6. Trading signals with confidence levels
 7. Risk management suggestions
 
-Format the response in markdown with clear sections, using professional trading terminology.
+Format response in markdown with clear sections using professional trading terminology.<|eot_id|><|start_header_id|>assistant<|end_header_id|>
 """
 
-# Chart Analysis Functions
+# Helper functions
+def image_to_base64(img):
+    buffered = io.BytesIO()
+    img.save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode("utf-8")
+
 def preprocess_image(image):
     img = np.array(image)
     img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-    height, width = img.shape[:2]
     
     # Resize if necessary
+    height, width = img.shape[:2]
     if width > MAX_WIDTH or height > MAX_HEIGHT:
         img = cv2.resize(img, (MAX_WIDTH, MAX_HEIGHT), interpolation=cv2.INTER_AREA)
     
@@ -60,14 +85,21 @@ def preprocess_image(image):
     lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
     l, a, b = cv2.split(lab)
     clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8,8))
-    limg = cv2.merge([clahe.apply(l),a,b])
+    limg = cv2.merge([clahe.apply(l), a, b])
     enhanced = cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)
     
-    return enhanced
+    return Image.fromarray(cv2.cvtColor(enhanced, cv2.COLOR_BGR2RGB))
 
 def detect_chart_elements(image):
-    # Detect technical patterns using YOLO model
-    detections = chart_detector(image)
+    inputs = detection_processor(images=image, return_tensors="pt").to(device)
+    outputs = detection_model(**inputs)
+    
+    target_sizes = torch.tensor([image.size[::-1]]).to(device)
+    results = detection_processor.post_process_object_detection(
+        outputs, 
+        target_sizes=target_sizes, 
+        threshold=0.8
+    )[0]
     
     elements = {
         'support_resistance': [],
@@ -76,65 +108,89 @@ def detect_chart_elements(image):
         'candlesticks': [],
     }
     
-    # Process model detections
-    for det in detections:
-        label = det['label']
-        if 'support' in label.lower() or 'resistance' in label.lower():
-            elements['support_resistance'].append(label)
-        elif 'trendline' in label.lower():
-            elements['trendlines'].append(label)
-        elif 'pattern' in label.lower():
-            elements['patterns'].append(label.split('_')[0])  # Assuming the pattern name is before '_'
-        elif 'candlestick' in label.lower():
-            elements['candlesticks'].append(label)
+    # Draw annotations and categorize elements
+    draw = ImageDraw.Draw(image)
+    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+        box = [round(i, 2) for i in box.tolist()]
+        label_name = detection_model.config.id2label[label.item()]
+        
+        # Draw bounding box
+        draw.rectangle(box, outline="red", width=2)
+        draw.text((box[0], box[1]), f"{label_name} ({round(score.item(), 2)})", fill="red")
+        
+        # Categorize elements (customize these mappings based on your detection model's labels)
+        if "support" in label_name.lower() or "resistance" in label_name.lower():
+            elements['support_resistance'].append(label_name)
+        elif "trendline" in label_name.lower():
+            elements['trendlines'].append(label_name)
+        elif "pattern" in label_name.lower():
+            elements['patterns'].append(label_name)
+        elif "candlestick" in label_name.lower():
+            elements['candlesticks'].append(label_name)
     
-    return elements
+    return image, elements
 
 def generate_llm_response(elements, question):
     prompt = SYSTEM_PROMPT.format(
-        support_resistance="; ".join(elements['support_resistance']),
-        trendlines="; ".join(elements['trendlines']),
-        patterns="; ".join(elements['patterns']),
-        candlesticks="; ".join(elements['candlesticks']),
+        support_resistance=", ".join(elements['support_resistance']),
+        trendlines=", ".join(elements['trendlines']),
+        patterns=", ".join(elements['patterns']),
+        candlesticks=", ".join(elements['candlesticks']),
         question=question
     )
     
-    inputs = tokenizer(prompt, return_tensors="pt", max_length=1024, truncation=True)
-    outputs = model.generate(
-        input_ids=inputs.input_ids,
-        attention_mask=inputs.attention_mask,
-        max_length=1500,
-        temperature=0.3,
-        num_beams=5,
-        early_stopping=True
+    inputs = llm_tokenizer(prompt, return_tensors="pt").to(device)
+    outputs = llm_model.generate(
+        **inputs,
+        max_new_tokens=1500,
+        temperature=0.7,
+        top_p=0.9,
+        do_sample=True,
+        pad_token_id=llm_tokenizer.eos_token_id
     )
     
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
+    response = llm_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response.split("<|start_header_id|>assistant<|end_header_id|>")[-1].strip()
 
-# Gradio Interface Functions
-def respond(image, question, chat_history):
-    # Preprocess image
-    processed_img = preprocess_image(image)
+# Gradio Interface
+def respond(message, history, image):
+    if image is None:
+        return "Please upload a chart image first."
     
-    # Detect elements in the chart
-    elements = detect_chart_elements(processed_img)
+    # Preprocess and analyze image
+    processed_img = preprocess_image(image)
+    annotated_img, elements = detect_chart_elements(processed_img)
     
-    # Generate analysis response
-    analysis = generate_llm_response(elements, question)
+    # Generate analysis
+    analysis = generate_llm_response(elements, message)
     
-    # Update chat history with the user's question and the model's analysis response
-    chat_history.append(("User: " + question, "Assistant: " + analysis))
+    # Convert annotated image to base64
+    img_base64 = image_to_base64(annotated_img)
+    img_html = f'<img src="data:image/png;base64,{img_base64}" style="max-width: 800px; margin-bottom: 20px;">'
     
-    return chat_history
+    return f"{img_html}\n{analysis}"
 
-# Gradio Interface
+# Create interface
 demo = gr.ChatInterface(
     fn=respond,
-    chatbot=gr.Chatbot(show_copy_button=True, layout="panel"),
-    title="Crypto Trading Assistant",
+    additional_inputs=[gr.Image(label="Upload Chart", type="pil")],
+    chatbot=gr.Chatbot(
+        show_copy_button=True,
+        layout="panel",
+        bubble_full_width=False,
+        sanitize_html=False
+    ),
+    title="Crypto Trading Assistant Pro",
     theme="Nymbo/Nymbo_Theme",
-    textbox=gr.Textbox(label="Ask Technical Questions", placeholder="Enter your analysis question...")
+    textbox=gr.Textbox(
+        label="Ask Technical Questions",
+        placeholder="Upload chart image and ask analysis questions...",
+        container=False
+    ),
+    examples=[
+        ["Is this a bullish reversal pattern?", "chart1.png"],
+        ["What are the key support levels?", "chart2.jpg"]
+    ]
 )
 
 if __name__ == "__main__":