Update app.py

app.py

@@ -130,7 +130,7 @@ import cv2
 import os
 
 # Load the model
-model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt')
+model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
 
 def detect_image(image):
     results = model(image)
@@ -145,12 +145,15 @@ def detect_video(video_path):
     os.makedirs(frames_dir, exist_ok=True)
 
     frame_count = 0
+    process_frame_count = 0
     while True:
         success, frame = video.read()
         if not success:
             break
-        frame_output_path = os.path.join(frames_dir, f'frame_{frame_count:04d}.jpg')
-        cv2.imwrite(frame_output_path, frame)
+        if frame_count % 5 == 0:  # Process every 5th frame (adjust as needed)
+            frame_output_path = os.path.join(frames_dir, f'frame_{process_frame_count:04d}.jpg')
+            cv2.imwrite(frame_output_path, frame)
+            process_frame_count += 1
         frame_count += 1
 
     video.release()
@@ -160,7 +163,7 @@ def detect_video(video_path):
     results_dir = 'results'
     os.makedirs(results_dir, exist_ok=True)
 
-    for i in range(frame_count):
+    for i in range(process_frame_count):
         frame_path = os.path.join(frames_dir, f'frame_{i:04d}.jpg')
         frame = cv2.imread(frame_path)
         results = model(frame)
@@ -175,7 +178,7 @@ def detect_video(video_path):
 
     video_output_path = 'output_video.mp4'  # Replace with your desired output video path
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # You can change the codec as needed
-    video_writer = cv2.VideoWriter(video_output_path, fourcc, frame_rate, (width, height))
+    video_writer = cv2.VideoWriter(video_output_path, fourcc, frame_rate/5, (width, height))  # Adjust frame rate
 
     for frame_file in frame_files:
         frame_path = os.path.join(results_dir, frame_file)
@@ -213,3 +216,4 @@ tabbed_interface = gr.TabbedInterface(interfaces, ["Image", "Video"])
 
 # Launch the tabbed interface 
 tabbed_interface.launch(debug=True)
+