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

Merge 3dd503c255825226da883314d3bc3a87c3de0480 into 453c6e38a51e9d1d5a2aa5fb7f1014a711913397

Browse Source
This commit is contained in:
cho-96 2025-03-26 06:50:33 +00:00 committed by GitHub
commit 475c4c7874
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 7 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
app.py
View File

@ -7,10 +7,12 @@ from ultralytics import YOLOv10
def yolov10_inference(image, video, model_id, image_size, conf_threshold): def yolov10_inference(image, video, model_id, image_size, conf_threshold):
model = YOLOv10.from_pretrained(f'jameslahm/{model_id}') model = YOLOv10.from_pretrained(f'jameslahm/{model_id}')
if image: if image:
assert image is not None, "Image input is required."
results = model.predict(source=image, imgsz=image_size, conf=conf_threshold) results = model.predict(source=image, imgsz=image_size, conf=conf_threshold)
annotated_image = results[0].plot() annotated_image = results[0].plot()
return annotated_image[:, :, ::-1], None return annotated_image[:, :, ::-1], None
else: else:
assert video is not None, "Video input is required."
video_path = tempfile.mktemp(suffix=".webm") video_path = tempfile.mktemp(suffix=".webm")
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:
@ -90,7 +92,7 @@ def app():
def update_visibility(input_type): def update_visibility(input_type):
image = gr.update(visible=True) if input_type == "Image" else gr.update(visible=False) image = gr.update(visible=True) if input_type == "Image" else gr.update(visible=False)
video = gr.update(visible=False) if input_type == "Image" else gr.update(visible=True) video = gr.update(visible=False) if input_type == "Image" else gr.update(visible=True)
output_image = gr.update(visible=True) if input_type == "Image" else gr.update(visible=False) output_image = gr.update(visible=True) if input_type == "Image" else gr.update(visible(False))
output_video = gr.update(visible=False) if input_type == "Image" else gr.update(visible=True) output_video = gr.update(visible=False) if input_type == "Image" else gr.update(visible=True)
return image, video, output_image, output_video return image, video, output_image, output_video
@ -102,6 +104,10 @@ 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):
assert model_id in ["yolov10n", "yolov10s", "yolov10m", "yolov10b", "yolov10l", "yolov10x"], "Invalid model ID."
assert 320 <= image_size <= 1280, "Image size out of range."
assert 0.0 <= conf_threshold <= 1.0, "Confidence threshold out of range."
if input_type == "Image": if input_type == "Image":
return yolov10_inference(image, None, model_id, image_size, conf_threshold) return yolov10_inference(image, None, model_id, image_size, conf_threshold)
else: else: