xclele/yolov10
1
0
Fork 0
mirror of https://github.com/THU-MIG/yolov10.git synced 2025-05-29 02:14:22 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

input video size is limited

Browse Source
This commit is contained in:
Sencer Yücel 2024-05-31 14:29:06 +03:00
parent 34af6d27fa
commit 1771746f62
1 changed files with 42 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

72
app.py
View File

@ -6,45 +6,57 @@ from ultralytics import YOLOv10
import cv2 import cv2
import tempfile 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) model = YOLOv10(model_path)
results = model.predict(source=image, imgsz=image_size, conf=conf_threshold) if image:
annotated_image = results[0].plot() results = model.predict(source=image, imgsz=image_size, conf=conf_threshold)
return annotated_image[:, :, ::-1], None 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): cap = cv2.VideoCapture(video_path)
model = YOLOv10(model_path) fps = cap.get(cv2.CAP_PROP_FPS)
video_path = tempfile.mktemp(suffix=".mp4") frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
with open(video_path, "wb") as f: frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
with open(video, "rb") as g:
f.write(g.read())
cap = cv2.VideoCapture(video_path) output_video_path = tempfile.mktemp(suffix=".mp4")
fps = cap.get(cv2.CAP_PROP_FPS) out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
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") while cap.isOpened():
out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height)) ret, frame = cap.read()
if not ret:
break
while cap.isOpened(): results = model.predict(source=frame, imgsz=image_size, conf=conf_threshold)
ret, frame = cap.read() annotated_frame = results[0].plot()
if not ret: out.write(annotated_frame)
break
results = model.predict(source=frame, imgsz=image_size, conf=conf_threshold) cap.release()
annotated_frame = results[0].plot() out.release()
out.write(annotated_frame)
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): 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 return annotated_image
def app(): def app():
with gr.Blocks(): with gr.Blocks():
with gr.Row(): with gr.Row():
@ -105,9 +117,9 @@ def app():
def run_inference(image, video, model_id, image_size, conf_threshold, input_type): def run_inference(image, video, model_id, image_size, conf_threshold, input_type):
if input_type == "Image": 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: 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( yolov10_infer.click(
fn=run_inference, fn=run_inference,