nehulagrawal commited on
Commit
31bbb9a
·
verified ·
1 Parent(s): 2721980

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +71 -53
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 = "foduucom/stockmarket-pattern-detection-yolov8"
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 time.time() - start_time < 60:
145
- # Capture the screen region
146
  sct_img = sct.grab(monitor)
147
  img = np.array(sct_img)
148
  img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
149
 
150
- # Save the frame in the screenshots folder
151
- timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
152
- image_name = f"predicted_images_{timestamp}_{frame_count}.png"
153
- image_path = os.path.join(screenshots_path, image_name)
154
- cv2.imwrite(image_path, img)
155
-
156
- # Run YOLO model and get save directory
157
- results = model(image_path, save=True)
158
- predict_path = results[0].save_dir if results else None
159
-
160
- # Find the latest annotated image inside predict_path
161
- if predict_path and os.path.exists(predict_path):
162
- annotated_images = sorted(glob.glob(os.path.join(predict_path, "*.jpg")), key=os.path.getmtime, reverse=True)
163
- final_image_path = annotated_images[0] if annotated_images else image_path
164
- else:
165
- final_image_path = image_path # Fallback to original image
166
-
167
- # Determine predicted label
168
- if results and results[0].boxes:
169
- class_indices = results[0].boxes.cls.tolist()
170
- predicted_label = classes[int(class_indices[0])]
171
- else:
172
- predicted_label = "No pattern detected"
173
-
174
- # Insert data into Excel (store path instead of image)
175
- ws.append([timestamp, final_image_path, predicted_label])
176
-
177
- # Read the image for video processing
178
- annotated_img = cv2.imread(final_image_path)
179
- if annotated_img is not None:
180
- # Add timestamp and label text to the image
181
- font = cv2.FONT_HERSHEY_SIMPLEX
182
- cv2.putText(annotated_img, f"{timestamp}", (10, 30), font, 0.7, (0, 255, 0), 2, cv2.LINE_AA)
183
- cv2.putText(annotated_img, f"{predicted_label}", (10, 60), font, 0.7, (0, 255, 255), 2, cv2.LINE_AA)
184
-
185
- # Initialize video writer if not already initialized
186
- if video_writer is None:
187
- height, width, layers = annotated_img.shape
188
- video_writer = cv2.VideoWriter(video_path, fourcc, fps, (width, height))
189
-
190
- video_writer.write(annotated_img)
191
-
192
- print(f"Frame {frame_count}: {final_image_path} -> {predicted_label}")
193
- frame_count += 1
194
- time.sleep(5)
195
-
196
- # Save the Excel file
 
 
 
 
 
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
  }