Views
No views yet
['Head and shoulders bottom', 'Head and shoulders top', 'M_Head', 'StockLine', 'Triangle', 'W_Bottom']pip install mss==10.0.0 opencv-python==4.11.0.86 numpy ultralytics==8.3.94 openpyxl==3.1.51
2
3import os
4import mss # type: ignore
5import cv2
6import numpy as np
7import time
8import glob
9from ultralytics import YOLO
10from openpyxl import Workbook
11
12# Get the user's home directory
13home_dir = os.path.expanduser("~")
14
15# Define dynamic paths
16save_path = os.path.join(home_dir, "yolo_detection")
17screenshots_path = os.path.join(save_path, "screenshots")
18detect_path = os.path.join(save_path, "runs", "detect")
19
20# Ensure necessary directories exist
21os.makedirs(screenshots_path, exist_ok=True)
22os.makedirs(detect_path, exist_ok=True)
23
24# Define pattern classes
25classes = ['Head and shoulders bottom', 'Head and shoulders top', 'M_Head', 'StockLine', 'Triangle', 'W_Bottom']
26
27# Load YOLOv8 model
28model_path = "model.pt"
29if not os.path.exists(model_path):
30 raise FileNotFoundError(f"Model file not found: {model_path}")
31model = YOLO(model_path)
32
33# Define screen capture region
34monitor = {"top": 0, "left": 683, "width": 683, "height": 768}
35
36# Create an Excel file
37excel_file = os.path.join(save_path, "classification_results.xlsx")
38wb = Workbook()
39ws = wb.active
40ws.append(["Timestamp", "Predicted Image Path", "Label"]) # Headers
41
42# Initialize video writer
43video_path = os.path.join(save_path, "annotated_video.mp4")
44fourcc = cv2.VideoWriter_fourcc(*"mp4v")
45fps = 0.5 # Adjust frames per second as needed
46video_writer = None
47
48with mss.mss() as sct:
49 start_time = time.time()
50 last_capture_time = start_time # Track the last capture time
51 frame_count = 0
52
53 while True:
54 # Continuously capture the screen
55 sct_img = sct.grab(monitor)
56 img = np.array(sct_img)
57 img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
58
59 # Check if 60 seconds have passed since last YOLO prediction
60 current_time = time.time()
61 if current_time - last_capture_time >= 60:
62 # Take screenshot for YOLO prediction
63 timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
64 image_name = f"predicted_images_{timestamp}_{frame_count}.png"
65 image_path = os.path.join(screenshots_path, image_name)
66 cv2.imwrite(image_path, img)
67
68 # Run YOLO model and get save directory
69 results = model(image_path, save=True)
70 predict_path = results[0].save_dir if results else None
71
72 # Find the latest annotated image inside predict_path
73 if predict_path and os.path.exists(predict_path):
74 annotated_images = sorted(glob.glob(os.path.join(predict_path, "*.jpg")), key=os.path.getmtime, reverse=True)
75 final_image_path = annotated_images[0] if annotated_images else image_path
76 else:
77 final_image_path = image_path # Fallback to original image
78
79 # Determine predicted label
80 if results and results[0].boxes:
81 class_indices = results[0].boxes.cls.tolist()
82 predicted_label = classes[int(class_indices[0])]
83 else:
84 predicted_label = "No pattern detected"
85
86 # Insert data into Excel (store path instead of image)
87 ws.append([timestamp, final_image_path, predicted_label])
88
89 # Read the image for video processing
90 annotated_img = cv2.imread(final_image_path)
91 if annotated_img is not None:
92 # Add timestamp and label text to the image
93 font = cv2.FONT_HERSHEY_SIMPLEX
94 cv2.putText(annotated_img, f"{timestamp}", (10, 30), font, 0.7, (0, 255, 0), 2, cv2.LINE_AA)
95 cv2.putText(annotated_img, f"{predicted_label}", (10, 60), font, 0.7, (0, 255, 255), 2, cv2.LINE_AA)
96
97 # Initialize video writer if not already initialized
98 if video_writer is None:
99 height, width, layers = annotated_img.shape
100 video_writer = cv2.VideoWriter(video_path, fourcc, fps, (width, height))
101
102 video_writer.write(annotated_img)
103
104 print(f"Frame {frame_count}: {final_image_path} -> {predicted_label}")
105 frame_count += 1
106
107 # Update the last capture time
108 last_capture_time = current_time
109
110 # Save the Excel file periodically
111 wb.save(excel_file)
112
113 # If you want to continuously display the screen, you can add this line
114 cv2.imshow("Screen Capture", img)
115
116 # Break if 'q' is pressed (you can exit the loop this way)
117 if cv2.waitKey(1) & 0xFF == ord('q'):
118 break
119
120# Release video writer
121if video_writer is not None:
122 video_writer.release()
123 print(f"Video saved at {video_path}")
124
125# Remove all files in screenshots directory
126for file in os.scandir(screenshots_path):
127 os.remove(file.path)
128os.rmdir(screenshots_path)
129
130print(f"Results saved to {excel_file}")
131
132# Close OpenCV window
133cv2.destroyAllWindows()
1341@ModelCard{
2 author = {Nehul Agrawal,
3 Pranjal Singh Thakur, Priyal Mehta and Arjun Singh},
4 title = {YOLOv8s Stock Market Pattern Detection from Live Screen Capture},
5 year = {2023}
6}