Spaces:
Sleeping
Sleeping
File size: 3,842 Bytes
23b9701 2f60924 23b9701 1d9fcca 2f60924 1d9fcca 2f60924 23b9701 2f60924 759794a 2f60924 759794a 2f60924 759794a 2f60924 759794a 2f60924 a3ced2c 2f60924 759794a 2f60924 b283eca 23b9701 2f60924 23b9701 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
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()
|