Spaces:
Build error
Build error
| 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 |