Update app.py

app.py

@@ -3,36 +3,29 @@ import pandas as pd
 import plotly.graph_objects as go
 from ultralytics import YOLO
 import cv2
-import time
+import os
 import gradio as gr
 
-API_KEY = "ITWJ6NDTF45CBTDO"
+API_KEY = "ITWJ6NDTF45CBTDO"  # Consider using environment variables for API keys
 
-def get_stock_candlestick_data(symbol, interval="5min", output_size="compact"):
-    """
-    Fetch stock candlestick data from Alpha Vantage.
-    """
+def get_stock_candlestick_data(symbol, interval="1min", output_size="compact"):
+    """Fetch stock candlestick data from Alpha Vantage."""
     url = f"https://www.alphavantage.co/query?function=TIME_SERIES_INTRADAY&symbol={symbol}&interval={interval}&apikey={API_KEY}&outputsize={output_size}"
-    print(f"Fetching data from: {url}")  # Debugging
     response = requests.get(url)
     if response.status_code == 200:
         data = response.json()
-        print("API Response:", data)  # Debugging
         if f"Time Series ({interval})" in data:
             return data[f"Time Series ({interval})"]
         else:
-            print("Error: No candlestick data found in response.")
-            print(data)
             return None
     else:
-        print(f"Error fetching data: {response.status_code}")
-        print(response.text)
         return None
 
 def process_stock_candlestick_data(data):
-    """
-    Process Alpha Vantage stock candlestick data into a DataFrame.
-    """
+    """Process Alpha Vantage stock candlestick data into a DataFrame."""
+    if not data:
+        return None
+        
     rows = []
     for timestamp, values in data.items():
         rows.append({
@@ -43,12 +36,15 @@ def process_stock_candlestick_data(data):
             "close": float(values["4. close"]),
             "volume": float(values["5. volume"])
         })
-    return pd.DataFrame(rows)
+    df = pd.DataFrame(rows)
+    df = df.sort_values("timestamp")  # Ensure chronological order
+    return df
 
-def generate_candlestick_chart(df, n=50):
-    """
-    Generate a candlestick chart using Plotly with the last n data points.
-    """
+def generate_candlestick_chart(df, n=50, output_path="candlestick.png"):
+    """Generate a candlestick chart using Plotly with the last n data points."""
+    if df is None or len(df) == 0:
+        return None
+        
     df = df.tail(n)  # Use only the last n rows
     fig = go.Figure(data=[go.Candlestick(
         x=df["timestamp"],
@@ -63,65 +59,89 @@ def generate_candlestick_chart(df, n=50):
         yaxis_title="Price",
         xaxis_rangeslider_visible=False
     )
-    # Removed fig.show() since Gradio will display the image
-    fig.write_image("candlestick.png")
+    fig.write_image(output_path)
+    return output_path
 
-def yolo_model(img_path, model):
-    """
-    Run YOLO model on the image and count GAP UP and GAP DOWN patterns.
-    """
-    results = model(img_path)
-    gap_up_count = 0
-    gap_down_count = 0
-    for result in results:
-        classes = result.boxes.cls
-        for cls in classes:
-            if cls == 0:
-                gap_down_count += 1
-            elif cls == 1:
-                gap_up_count += 1
-        annotated_image = result.plot()
-    return annotated_image, gap_up_count, gap_down_count
+def yolo_model(img_path, model_path):
+    """Run YOLO model on the image and count GAP UP and GAP DOWN patterns."""
+    if not os.path.exists(img_path):
+        return None, 0, 0
+        
+    # Load model each time to avoid persistence issues in Spaces
+    try:
+        model = YOLO(model_path)
+        results = model(img_path)
+        gap_up_count = 0
+        gap_down_count = 0
+        
+        for result in results:
+            boxes = result.boxes
+            if hasattr(boxes, 'cls') and len(boxes.cls) > 0:
+                classes = boxes.cls.cpu().numpy() if hasattr(boxes.cls, 'cpu') else boxes.cls
+                for cls in classes:
+                    if int(cls) == 0:
+                        gap_down_count += 1
+                    elif int(cls) == 1:
+                        gap_up_count += 1
+                        
+        annotated_image = results[0].plot()
+        output_path = "annotated_output.png"
+        cv2.imwrite(output_path, annotated_image)
+        return output_path, gap_up_count, gap_down_count
+    except Exception as e:
+        print(f"Error running YOLO model: {e}")
+        return None, 0, 0
 
-def detect_gap_patterns(symbol):
-    """
-    Main function to fetch data, generate charts, and detect GAP patterns in near-real-time.
-    """
-    model = YOLO("/best.pt")  # Load model once outside the loop
-    while True:
-        data = get_stock_candlestick_data(symbol)
-        if not data:
-            print("Failed to fetch data. Retrying in 15 seconds.")
-            time.sleep(15)
-            continue  # Retry instead of exiting
-
-        df = process_stock_candlestick_data(data)
-        generate_candlestick_chart(df, n=50)  # Generate chart with last 50 candles
-        annotated_image, gap_up_count, gap_down_count = yolo_model("candlestick.png", model)
-        cv2.imwrite("annotated_output.png", annotated_image)
-        yield "annotated_output.png", gap_up_count, gap_down_count
-        time.sleep(15)  # Wait 15 seconds to respect API rate limits
+def detect_gap_patterns(symbol, model_path="best.pt"):
+    """Non-streaming function to fetch data, generate charts, and detect GAP patterns."""
+    # Check if the model file exists
+    if not os.path.exists(model_path):
+        return None, f"Model not found at {model_path}", f"Model not found at {model_path}"
+    
+    # Get stock data
+    data = get_stock_candlestick_data(symbol)
+    if not data:
+        return None, "Failed to fetch stock data", "Failed to fetch stock data"
+    
+    # Process data and generate chart
+    df = process_stock_candlestick_data(data)
+    if df is None or len(df) == 0:
+        return None, "No valid stock data available", "No valid stock data available"
+    
+    chart_path = generate_candlestick_chart(df, n=50)
+    if not chart_path or not os.path.exists(chart_path):
+        return None, "Failed to generate chart", "Failed to generate chart"
+    
+    # Run YOLO detection
+    annotated_path, gap_up_count, gap_down_count = yolo_model(chart_path, model_path)
+    if not annotated_path:
+        return None, "Failed to run detection model", "Failed to run detection model"
+    
+    return annotated_path, f"GAP UP Count: {gap_up_count}", f"GAP DOWN Count: {gap_down_count}"
 
 # Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("# GAP Pattern Detection in Real-Time Stock Charts")
-    gr.Markdown("Enter a stock symbol (e.g., AAPL) to detect GAP UP and GAP DOWN patterns in near-real-time candlestick charts.")
-    
+    gr.Markdown("# GAP Pattern Detection in Stock Charts")
+    gr.Markdown("Enter a stock symbol (e.g., AAPL) to detect GAP UP and GAP DOWN patterns in candlestick charts.")
+
     with gr.Row():
         symbol_input = gr.Textbox(label="Stock Symbol", placeholder="Enter a stock symbol (e.g., AAPL)")
-        submit_button = gr.Button("Start Real-Time Detection")
-    
+        model_path_input = gr.Textbox(label="Model Path", value="best.pt", placeholder="Path to YOLO model file")
+        submit_button = gr.Button("Detect Patterns")
+
     with gr.Row():
         output_image = gr.Image(label="Annotated Candlestick Chart")
-        gap_up_output = gr.Textbox(label="GAP UP Count")
-        gap_down_output = gr.Textbox(label="GAP DOWN Count")
-    
-    # Start real-time detection when the button is clicked
+        
+    with gr.Row():
+        gap_up_output = gr.Textbox(label="GAP UP Results")
+        gap_down_output = gr.Textbox(label="GAP DOWN Results")
+
+    # Run detection when the button is clicked
     submit_button.click(
         fn=detect_gap_patterns,
-        inputs=symbol_input,
+        inputs=[symbol_input, model_path_input],
         outputs=[output_image, gap_up_output, gap_down_output]
     )
 
 # Launch the Gradio app
-demo.launch(share=True, debug=True)
+demo.launch()