xclele/yolov10
1
0
Fork 0
mirror of https://github.com/THU-MIG/yolov10.git synced 2025-07-17 12:45:40 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

container and codec formats changed for video input/outputs

Browse Source
This commit is contained in:
Sencer Yücel 2024-06-03 11:25:29 +03:00
parent 167ea34027
commit b565076dd1
1 changed files with 3 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
app.py
View File

@ -7,15 +7,6 @@ import cv2
import tempfile import tempfile
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): def yolov10_inference(image, video, model_path, image_size, conf_threshold):
model = YOLOv10(model_path) model = YOLOv10(model_path)
if image: if image:
@ -23,8 +14,7 @@ def yolov10_inference(image, video, model_path, image_size, conf_threshold):
annotated_image = results[0].plot() annotated_image = results[0].plot()
return annotated_image[:, :, ::-1], None return annotated_image[:, :, ::-1], None
else: else:
check_file_size(video) video_path = tempfile.mktemp(suffix=".webm")
video_path = tempfile.mktemp(suffix=".mp4")
with open(video_path, "wb") as f: with open(video_path, "wb") as f:
with open(video, "rb") as g: with open(video, "rb") as g:
f.write(g.read()) f.write(g.read())
@ -34,8 +24,8 @@ def yolov10_inference(image, video, model_path, image_size, conf_threshold):
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
output_video_path = tempfile.mktemp(suffix=".mp4") output_video_path = tempfile.mktemp(suffix=".webm")
out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height)) out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'vp80'), fps, (frame_width, frame_height))
while cap.isOpened(): while cap.isOpened():
ret, frame = cap.read() ret, frame = cap.read()