Spaces:
Build error
Build error
File size: 4,480 Bytes
d82ffcc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | import cv2
import numpy as np
import os
from ultralytics import YOLO
from typing import List, Dict, Any
from PIL import Image
from huggingface_hub import hf_hub_download
class PatternDetector:
def __init__(self, model_path: str = "foduucom/stockmarket-pattern-detection-yolov8"):
if "/" in model_path and not os.path.exists(model_path):
print(f"⬇️ Downloading model from Hugging Face: {model_path}")
try:
local_model_path = hf_hub_download(
repo_id=model_path,
filename="model.pt",
cache_dir="./models"
)
print(f"✅ Model downloaded to: {local_model_path}")
except Exception as e:
print(f"❌ Failed to download model: {e}")
print("⚠️ Falling back to default YOLOv8n model.")
local_model_path = "yolov8n.pt"
else:
local_model_path = model_path
self.model = YOLO(local_model_path)
self.supported_patterns = [
'Head and shoulders bottom', 'Head and shoulders top',
'M_Head', 'W_Bottom', 'Triangle', 'Flag', 'Pennant',
'Double Top', 'Double Bottom', 'Cup and Handle'
]
def detect_patterns(self, image: np.ndarray, conf_threshold: float = 0.3) -> List[Dict[str, Any]]:
if isinstance(image, Image.Image):
image = np.array(image)
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
results = self.model(image, conf=conf_threshold)
detections = []
for result in results:
boxes = result.boxes
if boxes is not None:
for box in boxes:
x1, y1, x2, y2 = box.xyxy[0].tolist()
conf = float(box.conf[0])
cls_id = int(box.cls[0])
pattern_name = self.model.names[cls_id]
detections.append({
'pattern': pattern_name,
'confidence': conf,
'bbox': (int(x1), int(y1), int(x2), int(y2)),
'center': ((x1 + x2) / 2, (y1 + y2) / 2)
})
return detections
def analyze_pattern_sentiment(self, detections: List[Dict]) -> Dict[str, Any]:
bullish_patterns = ['W_Bottom', 'Double Bottom', 'Cup and Handle', 'Head and shoulders bottom']
bearish_patterns = ['M_Head', 'Double Top', 'Head and shoulders top']
bullish_score = 0
bearish_score = 0
for det in detections:
pattern = det['pattern']
conf = det['confidence']
if pattern in bullish_patterns:
bullish_score += conf
elif pattern in bearish_patterns:
bearish_score += conf
total = bullish_score + bearish_score
sentiment = "NEUTRAL"
strength = 0.0
if total > 0:
if bullish_score > bearish_score:
sentiment = "BULLISH"
strength = bullish_score / total
else:
sentiment = "BEARISH"
strength = bearish_score / total
return {
'sentiment': sentiment,
'strength': strength,
'bullish_score': bullish_score,
'bearish_score': bearish_score,
'patterns_found': [d['pattern'] for d in detections],
'highest_confidence_pattern': max(detections, key=lambda x: x['confidence']) if detections else None
}
def draw_detections(self, image: np.ndarray, detections: List[Dict]) -> np.ndarray:
img_copy = image.copy()
for det in detections:
x1, y1, x2, y2 = det['bbox']
label = f"{det['pattern']}: {det['confidence']:.2f}"
bullish_patterns = ['W_Bottom', 'Double Bottom', 'Cup and Handle']
bearish_patterns = ['M_Head', 'Double Top']
if det['pattern'] in bullish_patterns:
color = (0, 255, 0)
elif det['pattern'] in bearish_patterns:
color = (0, 0, 255)
else:
color = (255, 255, 0)
cv2.rectangle(img_copy, (x1, y1), (x2, y2), color, 2)
cv2.putText(img_copy, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
return img_copy |