Update README.md

README.md

@@ -34,7 +34,7 @@ pipeline_tag: object-detection
   <img width="500" alt="foduucom/stockmarket-pattern-detection-yolov8" src="https://huggingface.co/foduucom/stockmarket-pattern-detection-yolov8/resolve/main/thumbnail.jpg">
 </div>
 
-# Model Card for YOLOv8s Stock Market Pattern Detection from Live Screen Capture
+# Model Card for YOLOv8s Stock Market Real Time Pattern Detection from Live Screen Capture
 
 ## Model Summary
 
@@ -84,15 +84,15 @@ Users should be aware of the model's limitations and potential biases. Testing a
 
 ## How to Get Started with the Model
 
-To begin using the YOLOv8s Stock Market Pattern Detection model, install the necessary libraries:
+To begin using the YOLOv8s Stock Market Real Time Pattern Detection model, install the necessary libraries:
 ```bash
-pip install opencv-python==4.11.0.86 numpy==2.1.3 mss==10.0.0 ultralytics==8.3.94 openpyxl==3.1.5
+pip install mss==10.0.0 opencv-python==4.11.0.86 numpy==2.1.3 ultralytics==8.3.94 openpyxl==3.1.5
 ```
 
 ### Screen Capture and Pattern Detection Implementation
 ```python
 import os
-import mss
+import mss # type: ignore
 import cv2
 import numpy as np
 import time
@@ -100,11 +100,15 @@ import glob
 from ultralytics import YOLO
 from openpyxl import Workbook
 
-# Define paths
+# Get the user's home directory
 home_dir = os.path.expanduser("~")
+
+# Define dynamic paths
 save_path = os.path.join(home_dir, "yolo_detection")
 screenshots_path = os.path.join(save_path, "screenshots")
 detect_path = os.path.join(save_path, "runs", "detect")
+
+# Ensure necessary directories exist
 os.makedirs(screenshots_path, exist_ok=True)
 os.makedirs(detect_path, exist_ok=True)
 
@@ -112,49 +116,96 @@ os.makedirs(detect_path, exist_ok=True)
 classes = ['Head and shoulders bottom', 'Head and shoulders top', 'M_Head', 'StockLine', 'Triangle', 'W_Bottom']
 
 # Load YOLOv8 model
-model = YOLO('foduucom/stockmarket-pattern-detection-yolov8')
+model_path = "foduucom/stockmarket-pattern-detection-yolov8"
+if not os.path.exists(model_path):
+    raise FileNotFoundError(f"Model file not found: {model_path}")
+model = YOLO(model_path)
 
 # Define screen capture region
 monitor = {"top": 0, "left": 683, "width": 683, "height": 768}
 
-# Create Excel file
+# Create an Excel file
 excel_file = os.path.join(save_path, "classification_results.xlsx")
 wb = Workbook()
 ws = wb.active
-ws.append(["Timestamp", "Predicted Image Path", "Label"])
+ws.append(["Timestamp", "Predicted Image Path", "Label"])  # Headers
 
 # Initialize video writer
 video_path = os.path.join(save_path, "annotated_video.mp4")
 fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-fps = 0.5
+fps = 0.5  # Adjust frames per second as needed
 video_writer = None
 
 # Start capturing
 with mss.mss() as sct:
     start_time = time.time()
     frame_count = 0
-    
+
     while time.time() - start_time < 60:
+        # Capture the screen region
         sct_img = sct.grab(monitor)
         img = np.array(sct_img)
         img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
+
+        # Save the frame in the screenshots folder
         timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
         image_name = f"predicted_images_{timestamp}_{frame_count}.png"
         image_path = os.path.join(screenshots_path, image_name)
         cv2.imwrite(image_path, img)
-        
+
+        # Run YOLO model and get save directory
         results = model(image_path, save=True)
         predict_path = results[0].save_dir if results else None
-        annotated_images = sorted(glob.glob(os.path.join(predict_path, "*.jpg")), key=os.path.getmtime, reverse=True) if predict_path else []
-        final_image_path = annotated_images[0] if annotated_images else image_path
-        
-        predicted_label = classes[int(results[0].boxes.cls.tolist()[0])] if results and results[0].boxes else "No pattern detected"
+
+        # Find the latest annotated image inside predict_path
+        if predict_path and os.path.exists(predict_path):
+            annotated_images = sorted(glob.glob(os.path.join(predict_path, "*.jpg")), key=os.path.getmtime, reverse=True)
+            final_image_path = annotated_images[0] if annotated_images else image_path
+        else:
+            final_image_path = image_path  # Fallback to original image
+
+        # Determine predicted label
+        if results and results[0].boxes:
+            class_indices = results[0].boxes.cls.tolist()
+            predicted_label = classes[int(class_indices[0])]
+        else:
+            predicted_label = "No pattern detected"
+
+        # Insert data into Excel (store path instead of image)
         ws.append([timestamp, final_image_path, predicted_label])
-        wb.save(excel_file)
-        
+
+        # Read the image for video processing
+        annotated_img = cv2.imread(final_image_path)
+        if annotated_img is not None:
+            # Add timestamp and label text to the image
+            font = cv2.FONT_HERSHEY_SIMPLEX
+            cv2.putText(annotated_img, f"{timestamp}", (10, 30), font, 0.7, (0, 255, 0), 2, cv2.LINE_AA)
+            cv2.putText(annotated_img, f"{predicted_label}", (10, 60), font, 0.7, (0, 255, 255), 2, cv2.LINE_AA)
+            
+            # Initialize video writer if not already initialized
+            if video_writer is None:
+                height, width, layers = annotated_img.shape
+                video_writer = cv2.VideoWriter(video_path, fourcc, fps, (width, height))
+            
+            video_writer.write(annotated_img)
+
+        print(f"Frame {frame_count}: {final_image_path} -> {predicted_label}")
         frame_count += 1
         time.sleep(5)
-    
+
+        # Save the Excel file
+        wb.save(excel_file)
+
+# Release video writer
+if video_writer is not None:
+    video_writer.release()
+    print(f"Video saved at {video_path}")
+
+# Remove all files in screenshots directory
+for file in os.scandir(screenshots_path):
+    os.remove(file.path)
+os.rmdir(screenshots_path)
+
 print(f"Results saved to {excel_file}")
 ```
 
@@ -164,7 +215,7 @@ For inquiries and contributions, please contact us at [email protected].
 ```bibtex
 @ModelCard{
   author    = {Nehul Agrawal,
-               Pranjal Singh Thakur, Arjun Singh},
+               Pranjal Singh Thakur, and Arjun Singh},
   title     = {YOLOv8s Stock Market Pattern Detection from Live Screen Capture},
   year      = {2023}
 }