Update main.py
Browse files
main.py
CHANGED
@@ -9,7 +9,7 @@ from ultralytics import YOLO
|
|
9 |
from openpyxl import Workbook
|
10 |
|
11 |
# Ensure necessary directories exist
|
12 |
-
save_path = "
|
13 |
screenshots_path = os.path.join(save_path, "screenshots")
|
14 |
detect_path = os.path.join(save_path, "runs/detect/")
|
15 |
|
@@ -20,7 +20,7 @@ os.makedirs(screenshots_path, exist_ok=True)
|
|
20 |
classes = ['Head and shoulders bottom', 'Head and shoulders top', 'M_Head', 'StockLine', 'Triangle', 'W_Bottom']
|
21 |
|
22 |
# Load YOLOv8 model
|
23 |
-
model_path = "
|
24 |
if not os.path.exists(model_path):
|
25 |
raise FileNotFoundError(f"Model file not found: {model_path}")
|
26 |
model = YOLO(model_path)
|
@@ -35,7 +35,7 @@ ws = wb.active
|
|
35 |
ws.append(["Timestamp", "Predicted Image Path", "Label"]) # Headers
|
36 |
|
37 |
# Initialize video writer
|
38 |
-
video_path =
|
39 |
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
|
40 |
fps = 0.5 # Adjust frames per second as needed
|
41 |
video_writer = None
|
|
|
9 |
from openpyxl import Workbook
|
10 |
|
11 |
# Ensure necessary directories exist
|
12 |
+
save_path = "./"
|
13 |
screenshots_path = os.path.join(save_path, "screenshots")
|
14 |
detect_path = os.path.join(save_path, "runs/detect/")
|
15 |
|
|
|
20 |
classes = ['Head and shoulders bottom', 'Head and shoulders top', 'M_Head', 'StockLine', 'Triangle', 'W_Bottom']
|
21 |
|
22 |
# Load YOLOv8 model
|
23 |
+
model_path = "model.pt"
|
24 |
if not os.path.exists(model_path):
|
25 |
raise FileNotFoundError(f"Model file not found: {model_path}")
|
26 |
model = YOLO(model_path)
|
|
|
35 |
ws.append(["Timestamp", "Predicted Image Path", "Label"]) # Headers
|
36 |
|
37 |
# Initialize video writer
|
38 |
+
video_path = "./videos/annotated_video.mp4"
|
39 |
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
|
40 |
fps = 0.5 # Adjust frames per second as needed
|
41 |
video_writer = None
|