Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import cv2 | |
| import os | |
| import tempfile | |
| import shutil | |
| from ultralytics import YOLO | |
| # 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 (70MB) | |
| MAX_FILE_SIZE_MB = 70 # Limit is set to 70MB | |
| 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 {MAX_FILE_SIZE_MB}MB size limit. Please upload a smaller video.", [] | |
| print("Video size is ", video_size/1024) | |
| # 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)) | |
| 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 = [] | |
| 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.3, save=False, 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 | |
| if extracted_frames: | |
| return f"Success: {len(extracted_frames)} frames extracted and processed.", extracted_frames | |
| else: | |
| return "No damage found on the road.", [] | |
| # 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.Textbox(label="Status"), # Display status or error message | |
| gr.Gallery(label="Output Images Results") # Display extracted frames in a gallery | |
| ], | |
| title="Road Damage Detection", # Interface title | |
| description="Upload a video with a maximum size of 70MB." | |
| ) | |
| # Launch the Gradio app | |
| iface.launch() | |