Update README.md
Browse files
README.md
CHANGED
@@ -91,6 +91,8 @@ pip install mss==10.0.0 opencv-python==4.11.0.86 numpy ultralytics==8.3.94 openp
|
|
91 |
|
92 |
### Screen Capture and Pattern Detection Implementation
|
93 |
```python
|
|
|
|
|
94 |
import os
|
95 |
import mss # type: ignore
|
96 |
import cv2
|
@@ -116,7 +118,7 @@ os.makedirs(detect_path, exist_ok=True)
|
|
116 |
classes = ['Head and shoulders bottom', 'Head and shoulders top', 'M_Head', 'StockLine', 'Triangle', 'W_Bottom']
|
117 |
|
118 |
# Load YOLOv8 model
|
119 |
-
model_path = "
|
120 |
if not os.path.exists(model_path):
|
121 |
raise FileNotFoundError(f"Model file not found: {model_path}")
|
122 |
model = YOLO(model_path)
|
@@ -136,66 +138,78 @@ fourcc = cv2.VideoWriter_fourcc(*"mp4v")
|
|
136 |
fps = 0.5 # Adjust frames per second as needed
|
137 |
video_writer = None
|
138 |
|
139 |
-
# Start capturing
|
140 |
with mss.mss() as sct:
|
141 |
start_time = time.time()
|
|
|
142 |
frame_count = 0
|
143 |
-
|
144 |
-
while
|
145 |
-
#
|
146 |
sct_img = sct.grab(monitor)
|
147 |
img = np.array(sct_img)
|
148 |
img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
|
149 |
|
150 |
-
#
|
151 |
-
|
152 |
-
|
153 |
-
|
154 |
-
|
155 |
-
|
156 |
-
|
157 |
-
|
158 |
-
|
159 |
-
|
160 |
-
|
161 |
-
|
162 |
-
|
163 |
-
|
164 |
-
|
165 |
-
|
166 |
-
|
167 |
-
|
168 |
-
|
169 |
-
|
170 |
-
|
171 |
-
|
172 |
-
|
173 |
-
|
174 |
-
|
175 |
-
|
176 |
-
|
177 |
-
|
178 |
-
|
179 |
-
|
180 |
-
#
|
181 |
-
|
182 |
-
|
183 |
-
|
184 |
-
|
185 |
-
|
186 |
-
|
187 |
-
|
188 |
-
|
189 |
-
|
190 |
-
|
191 |
-
|
192 |
-
|
193 |
-
|
194 |
-
|
195 |
-
|
196 |
-
|
|
|
|
|
|
|
|
|
|
|
197 |
wb.save(excel_file)
|
198 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
199 |
# Release video writer
|
200 |
if video_writer is not None:
|
201 |
video_writer.release()
|
@@ -207,6 +221,10 @@ for file in os.scandir(screenshots_path):
|
|
207 |
os.rmdir(screenshots_path)
|
208 |
|
209 |
print(f"Results saved to {excel_file}")
|
|
|
|
|
|
|
|
|
210 |
```
|
211 |
|
212 |
## Model Contact
|
@@ -215,7 +233,7 @@ For inquiries and contributions, please contact us at info@foduu.com.
|
|
215 |
```bibtex
|
216 |
@ModelCard{
|
217 |
author = {Nehul Agrawal,
|
218 |
-
Pranjal Singh Thakur, and Arjun Singh},
|
219 |
title = {YOLOv8s Stock Market Pattern Detection from Live Screen Capture},
|
220 |
year = {2023}
|
221 |
}
|
|
|
91 |
|
92 |
### Screen Capture and Pattern Detection Implementation
|
93 |
```python
|
94 |
+
|
95 |
+
|
96 |
import os
|
97 |
import mss # type: ignore
|
98 |
import cv2
|
|
|
118 |
classes = ['Head and shoulders bottom', 'Head and shoulders top', 'M_Head', 'StockLine', 'Triangle', 'W_Bottom']
|
119 |
|
120 |
# Load YOLOv8 model
|
121 |
+
model_path = "model.pt"
|
122 |
if not os.path.exists(model_path):
|
123 |
raise FileNotFoundError(f"Model file not found: {model_path}")
|
124 |
model = YOLO(model_path)
|
|
|
138 |
fps = 0.5 # Adjust frames per second as needed
|
139 |
video_writer = None
|
140 |
|
|
|
141 |
with mss.mss() as sct:
|
142 |
start_time = time.time()
|
143 |
+
last_capture_time = start_time # Track the last capture time
|
144 |
frame_count = 0
|
145 |
+
|
146 |
+
while True:
|
147 |
+
# Continuously capture the screen
|
148 |
sct_img = sct.grab(monitor)
|
149 |
img = np.array(sct_img)
|
150 |
img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
|
151 |
|
152 |
+
# Check if 60 seconds have passed since last YOLO prediction
|
153 |
+
current_time = time.time()
|
154 |
+
if current_time - last_capture_time >= 60:
|
155 |
+
# Take screenshot for YOLO prediction
|
156 |
+
timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
|
157 |
+
image_name = f"predicted_images_{timestamp}_{frame_count}.png"
|
158 |
+
image_path = os.path.join(screenshots_path, image_name)
|
159 |
+
cv2.imwrite(image_path, img)
|
160 |
+
|
161 |
+
# Run YOLO model and get save directory
|
162 |
+
results = model(image_path, save=True)
|
163 |
+
predict_path = results[0].save_dir if results else None
|
164 |
+
|
165 |
+
# Find the latest annotated image inside predict_path
|
166 |
+
if predict_path and os.path.exists(predict_path):
|
167 |
+
annotated_images = sorted(glob.glob(os.path.join(predict_path, "*.jpg")), key=os.path.getmtime, reverse=True)
|
168 |
+
final_image_path = annotated_images[0] if annotated_images else image_path
|
169 |
+
else:
|
170 |
+
final_image_path = image_path # Fallback to original image
|
171 |
+
|
172 |
+
# Determine predicted label
|
173 |
+
if results and results[0].boxes:
|
174 |
+
class_indices = results[0].boxes.cls.tolist()
|
175 |
+
predicted_label = classes[int(class_indices[0])]
|
176 |
+
else:
|
177 |
+
predicted_label = "No pattern detected"
|
178 |
+
|
179 |
+
# Insert data into Excel (store path instead of image)
|
180 |
+
ws.append([timestamp, final_image_path, predicted_label])
|
181 |
+
|
182 |
+
# Read the image for video processing
|
183 |
+
annotated_img = cv2.imread(final_image_path)
|
184 |
+
if annotated_img is not None:
|
185 |
+
# Add timestamp and label text to the image
|
186 |
+
font = cv2.FONT_HERSHEY_SIMPLEX
|
187 |
+
cv2.putText(annotated_img, f"{timestamp}", (10, 30), font, 0.7, (0, 255, 0), 2, cv2.LINE_AA)
|
188 |
+
cv2.putText(annotated_img, f"{predicted_label}", (10, 60), font, 0.7, (0, 255, 255), 2, cv2.LINE_AA)
|
189 |
+
|
190 |
+
# Initialize video writer if not already initialized
|
191 |
+
if video_writer is None:
|
192 |
+
height, width, layers = annotated_img.shape
|
193 |
+
video_writer = cv2.VideoWriter(video_path, fourcc, fps, (width, height))
|
194 |
+
|
195 |
+
video_writer.write(annotated_img)
|
196 |
+
|
197 |
+
print(f"Frame {frame_count}: {final_image_path} -> {predicted_label}")
|
198 |
+
frame_count += 1
|
199 |
+
|
200 |
+
# Update the last capture time
|
201 |
+
last_capture_time = current_time
|
202 |
+
|
203 |
+
# Save the Excel file periodically
|
204 |
wb.save(excel_file)
|
205 |
|
206 |
+
# If you want to continuously display the screen, you can add this line
|
207 |
+
cv2.imshow("Screen Capture", img)
|
208 |
+
|
209 |
+
# Break if 'q' is pressed (you can exit the loop this way)
|
210 |
+
if cv2.waitKey(1) & 0xFF == ord('q'):
|
211 |
+
break
|
212 |
+
|
213 |
# Release video writer
|
214 |
if video_writer is not None:
|
215 |
video_writer.release()
|
|
|
221 |
os.rmdir(screenshots_path)
|
222 |
|
223 |
print(f"Results saved to {excel_file}")
|
224 |
+
|
225 |
+
# Close OpenCV window
|
226 |
+
cv2.destroyAllWindows()
|
227 |
+
|
228 |
```
|
229 |
|
230 |
## Model Contact
|
|
|
233 |
```bibtex
|
234 |
@ModelCard{
|
235 |
author = {Nehul Agrawal,
|
236 |
+
Pranjal Singh Thakur, Priyal Mehta and Arjun Singh},
|
237 |
title = {YOLOv8s Stock Market Pattern Detection from Live Screen Capture},
|
238 |
year = {2023}
|
239 |
}
|