candlestick-api / app.py
doinglean's picture
Update app.py
d2e2940 verified
raw
history blame
1.63 kB
import gradio as gr
from huggingface_hub import hf_hub_download
from ultralytics import YOLO
import torch
# Lade das Modell
model_path = hf_hub_download(repo_id="foduucom/stockmarket-pattern-detection-yolov8", filename="model.pt")
model = YOLO(model_path)
def analyze_image(image, prompt):
# Verwende den Prompt (falls nötig, hier als Kontext für die Verarbeitung)
# YOLOv8 ignoriert den Prompt direkt, daher speichern wir ihn für die Logik
results = model.predict(source=image, save=False)
detections = []
for result in results:
for box in result.boxes:
label = result.names[int(box.cls)]
confidence = float(box.conf)
# Farben basierend auf Label oder Prompt (z. B. "bullish" für grün)
color = "green" if "bullish" in prompt.lower() or "Bullish" in label else "red"
detections.append({
"pattern": label,
"confidence": confidence,
"color": color,
"prompt_used": prompt # Rückgabe des Prompts zur Überprüfung
})
return detections
# Erstelle Gradio-Schnittstelle mit Bild- und Text-Eingabe
iface = gr.Interface(
fn=analyze_image,
inputs=[
gr.Image(type="pil", label="Upload TradingView Screenshot"),
gr.Textbox(label="Prompt", placeholder="Enter your prompt, e.g., 'Detect candlestick patterns and colors'")
],
outputs="json",
title="Candlestick Pattern Detection",
description="Upload a TradingView screenshot and provide a prompt to detect candlestick patterns and colors."
)
# Starte die App
iface.launch()