| import gradio as gr | |
| import PIL.Image as Image | |
| from ultralytics import YOLO | |
| classify = YOLO("models/classify.pt") | |
| def predict_image(image, conf_threshold, iou_threshold): | |
| results = classify.predict( | |
| image, conf=conf_threshold, iou=iou_threshold, stream=True) | |
| for r in results: | |
| im_array = r.plot(labels=True, boxes=True) | |
| yield Image.fromarray(im_array[..., ::-1]) | |
| iface = gr.Interface( | |
| fn=predict_image, | |
| inputs=[ | |
| gr.Video(label="Upload Video"), | |
| gr.Slider(minimum=0, maximum=1, value=0.85, | |
| label="Confidence threshold"), | |
| gr.Slider(minimum=0, maximum=1, value=0.7, label="IoU threshold"), | |
| ], | |
| outputs=gr.Image(type="numpy", label="Result"), | |
| title="Basketball Classifier", | |
| description="Have you ever wondered where the ball was when you were playing basketball? Where the rim was? Where you were? Videos may take a LOT of time since this is running on the basic CPU tier of HuggingFace. Feel free to check out the image space for a much faster demo!", | |
| ) | |
| if __name__ == "__main__": | |
| iface.launch() | |