video size added upto 50

app.py

@@ -4,59 +4,87 @@ import os
 import tempfile
 
 from ultralytics import YOLO
-model = YOLO(r"best.pt")  # initialize model
+
+# Initialize YOLO model
+model_path = r"best.pt"
+if not os.path.exists(model_path):
+    raise FileNotFoundError(f"Model file not found at {model_path}")
+model = YOLO(model_path)
+
+# Define the maximum file size (100MB)
+MAX_FILE_SIZE_MB = 50  # Limit is set to 100MB
+MAX_FILE_SIZE_BYTES = MAX_FILE_SIZE_MB * 1024 * 1024  # Convert MB to bytes
+
 
 
 def process_video(video_file):
+    # Check the size of the uploaded video
+    video_size = os.path.getsize(video_file)
+    if video_size > MAX_FILE_SIZE_BYTES:
+        return f"Error: The uploaded video exceeds the 100MB size limit. Please upload a smaller video."
+
     # Create a temporary folder to save extracted frames
     temp_dir = tempfile.mkdtemp()
-    
+
     # Capture the video using OpenCV
     video = cv2.VideoCapture(video_file)
+    if not video.isOpened():
+        return f"Error: Unable to open video file {video_file}"
+    
     fps = int(video.get(cv2.CAP_PROP_FPS))
-    frame_rate = fps // 5  # Extract 5 frames per second
+    if fps == 0:
+        return "Error: Video FPS is 0 or not detected."
+    
+    frame_rate = max(1, fps // 5)  # Extract 5 frames per second, ensure at least 1 frame/sec
 
     frame_count = 0
     extracted_frames = []
 
-    while True:
-        ret, frame = video.read()
-        if not ret:
-            break
-
-        # Extract every nth frame according to the frame_rate
-        if frame_count % frame_rate == 0:
-            # Resize the frame if it's larger than 1280x720
-            h, w, _ = frame.shape
-            if w > 1280 or h > 720:
-                frame = cv2.resize(frame, (1280, 720))
-
-            # Save the frame to the temporary directory
-            frame_path = os.path.join(temp_dir, f"frame_{frame_count}.jpg")
-            cv2.imwrite(frame_path, frame)
-            results = model.predict(source=frame_path, conf=0.2, save=False, show=False, line_width=2, show_labels=True, show_boxes=True, retina_masks=True)
-            for result in results:
-                boxes = result.boxes
-                # result.save(filename=frame_path)
-                # Check if any detections (boxes or masks, keypoints, etc.) exist
-                if boxes is not None and len(boxes) > 0:  # If boxes are present
-                    # Save the result image if detections are found
-                    result.save(filename=frame_path)
-                else:
-                    # If no detections, ensure the file doesn't exist
-                    if os.path.exists(frame_path):
-                        os.remove(frame_path)  # Delete the image if it exists
-
-            # Add the frame path to the list of extracted frames
-            if os.path.exists(frame_path):
-                extracted_frames.append(frame_path)
-
-        frame_count += 1
-
-    video.release()
+    try:
+        while True:
+            ret, frame = video.read()
+            if not ret:
+                break
+
+            # Extract every nth frame according to the frame_rate
+            if frame_count % frame_rate == 0:
+                # Resize the frame if it's larger than 1280x720
+                h, w, _ = frame.shape
+                if w > 1280 or h > 720:
+                    frame = cv2.resize(frame, (1280, 720))
+
+                # Save the frame to the temporary directory
+                frame_path = os.path.join(temp_dir, f"frame_{frame_count}.jpg")
+                cv2.imwrite(frame_path, frame)
+
+                # Perform detection using the YOLO model
+                results = model.predict(source=frame_path, conf=0.002, save=True, show=False, 
+                                        line_width=2, show_labels=True, show_boxes=True, retina_masks=True)
+                
+                for result in results:
+                    boxes = result.boxes
+                    if boxes is not None and len(boxes) > 0:
+                        # Save the result image if detections are found
+                        result.save(filename=frame_path)
+                    else:
+                        # If no detections, delete the frame
+                        if os.path.exists(frame_path):
+                            os.remove(frame_path)
+
+                # Add the frame path to the list of extracted frames if it still exists
+                if os.path.exists(frame_path):
+                    extracted_frames.append(frame_path)
+
+            frame_count += 1
+    finally:
+        video.release()
+        # Clean up temp directory after processing
+        if len(extracted_frames) == 0:
+            shutil.rmtree(temp_dir)
 
     # Return the list of extracted frame paths for display
-    return extracted_frames
+    return extracted_frames if extracted_frames else ["No frames extracted or no detections found."]
+
 
 # Define the Gradio interface
 iface = gr.Interface(

main.py

@@ -1,71 +0,0 @@
-import gradio as gr
-import cv2
-import os
-import tempfile
-
-from ultralytics import YOLO
-model = YOLO(r"best.pt")  # initialize model
-
-
-def process_video(video_file):
-    # Create a temporary folder to save extracted frames
-    temp_dir = tempfile.mkdtemp()
-    
-    # Capture the video using OpenCV
-    video = cv2.VideoCapture(video_file)
-    fps = int(video.get(cv2.CAP_PROP_FPS))
-    frame_rate = fps // 5  # Extract 5 frames per second
-
-    frame_count = 0
-    extracted_frames = []
-
-    while True:
-        ret, frame = video.read()
-        if not ret:
-            break
-
-        # Extract every nth frame according to the frame_rate
-        if frame_count % frame_rate == 0:
-            # Resize the frame if it's larger than 1280x720
-            h, w, _ = frame.shape
-            if w > 1280 or h > 720:
-                frame = cv2.resize(frame, (1280, 720))
-
-            # Save the frame to the temporary directory
-            frame_path = os.path.join(temp_dir, f"frame_{frame_count}.jpg")
-            cv2.imwrite(frame_path, frame)
-            results = model.predict(source=frame_path, conf=0.002, save=False, show=False, line_width=2, show_labels=True, show_boxes=True, retina_masks=True)
-            for result in results:
-                boxes = result.boxes
-                # result.save(filename=frame_path)
-                # Check if any detections (boxes or masks, keypoints, etc.) exist
-                if boxes is not None and len(boxes) > 0:  # If boxes are present
-                    # Save the result image if detections are found
-                    result.save(filename=frame_path)
-                else:
-                    # If no detections, ensure the file doesn't exist
-                    if os.path.exists(frame_path):
-                        os.remove(frame_path)  # Delete the image if it exists
-
-            # Add the frame path to the list of extracted frames
-            if os.path.exists(frame_path):
-                extracted_frames.append(frame_path)
-
-        frame_count += 1
-
-    video.release()
-
-    # Return the list of extracted frame paths for display
-    return extracted_frames
-
-# Define the Gradio interface
-iface = gr.Interface(
-    fn=process_video,  # Function to process the video and return extracted frames
-    inputs=gr.Video(label="Upload your video"),  # Video upload input
-    outputs=gr.Gallery(label="Extracted Frames"),  # Display extracted frames in a gallery
-    title="Video Frame Extraction",  # Interface title
-    description="Upload a video, extract 5 frames per second, resize if necessary, and display extracted frames."
-)
-
-# Launch the Gradio app
-iface.launch()