From 1771746f6227996062a74cb78bdb80ac82f24516 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sencer=20Y=C3=BCcel?= Date: Fri, 31 May 2024 14:29:06 +0300 Subject: [PATCH] input video size is limited --- app.py | 72 ++++++++++++++++++++++++++++++++++------------------------ 1 file changed, 42 insertions(+), 30 deletions(-) diff --git a/app.py b/app.py index 6a51e352..b5bc33d5 100644 --- a/app.py +++ b/app.py @@ -6,45 +6,57 @@ from ultralytics import YOLOv10 import cv2 import tempfile -def yolov10_inference(image, model_path, image_size, conf_threshold): + +def check_file_size(file, max_file_size=20*1024*1024): # 20MB + file_size = 0 + if file: + with open(file, "rb") as f: + file_size = len(f.read()) + if file_size > max_file_size: + raise gr.Error("File size exceeds the 20MB limit. Please try with another file.") + + +def yolov10_inference(image, video, model_path, image_size, conf_threshold): model = YOLOv10(model_path) - results = model.predict(source=image, imgsz=image_size, conf=conf_threshold) - annotated_image = results[0].plot() - return annotated_image[:, :, ::-1], None + if image: + results = model.predict(source=image, imgsz=image_size, conf=conf_threshold) + annotated_image = results[0].plot() + return annotated_image[:, :, ::-1], None + else: + check_file_size(video) + video_path = tempfile.mktemp(suffix=".mp4") + with open(video_path, "wb") as f: + with open(video, "rb") as g: + f.write(g.read()) -def yolov10_inference_video(video, model_path, image_size, conf_threshold): - model = YOLOv10(model_path) - video_path = tempfile.mktemp(suffix=".mp4") - with open(video_path, "wb") as f: - with open(video, "rb") as g: - f.write(g.read()) + cap = cv2.VideoCapture(video_path) + fps = cap.get(cv2.CAP_PROP_FPS) + frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) - cap = cv2.VideoCapture(video_path) - fps = cap.get(cv2.CAP_PROP_FPS) - frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) - frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + output_video_path = tempfile.mktemp(suffix=".mp4") + out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height)) - output_video_path = tempfile.mktemp(suffix=".mp4") - out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height)) + while cap.isOpened(): + ret, frame = cap.read() + if not ret: + break - while cap.isOpened(): - ret, frame = cap.read() - if not ret: - break + results = model.predict(source=frame, imgsz=image_size, conf=conf_threshold) + annotated_frame = results[0].plot() + out.write(annotated_frame) - results = model.predict(source=frame, imgsz=image_size, conf=conf_threshold) - annotated_frame = results[0].plot() - out.write(annotated_frame) + cap.release() + out.release() + + return None, output_video_path - cap.release() - out.release() - - return None, output_video_path def yolov10_inference_for_examples(image, model_path, image_size, conf_threshold): - annotated_image, _ = yolov10_inference(image, model_path, image_size, conf_threshold) + annotated_image, _ = yolov10_inference(image, None, model_path, image_size, conf_threshold) return annotated_image + def app(): with gr.Blocks(): with gr.Row(): @@ -105,9 +117,9 @@ def app(): def run_inference(image, video, model_id, image_size, conf_threshold, input_type): if input_type == "Image": - return yolov10_inference(image, model_id, image_size, conf_threshold) + return yolov10_inference(image, None, model_id, image_size, conf_threshold) else: - return yolov10_inference_video(video, model_id, image_size, conf_threshold) + return yolov10_inference(None, video, model_id, image_size, conf_threshold) yolov10_infer.click( fn=run_inference,