nehulagrawal commited on
Commit
29a606c
·
verified ·
1 Parent(s): 0b3f0b4

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +3 -3
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 = "/home/ml/ML/ml_backup/arjun/"
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 = "/home/ml/ML/ml_backup/arjun/best111.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,7 +35,7 @@ ws = wb.active
35
  ws.append(["Timestamp", "Predicted Image Path", "Label"]) # Headers
36
 
37
  # Initialize video writer
38
- video_path = os.path.join(save_path, "annotated_video.mp4")
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