| import gradio as gr |
| import torch |
| from sahi.prediction import ObjectPrediction |
| from sahi.utils.cv import visualize_object_predictions, read_image |
| from ultralyticsplus import YOLO |
|
|
|
|
|
|
| def yolov8_go( |
| image: gr.inputs.Image = None, |
| model_path: gr.inputs.Dropdown = None, |
| image_size: gr.inputs.Slider = 640, |
| conf_threshold: gr.inputs.Slider = 0.25, |
| iou_threshold: gr.inputs.Slider = 0.45, |
| ): |
|
|
| model = YOLO(model_path) |
| model.conf = conf_threshold |
| model.iou = iou_threshold |
| results = model.predict(image, imgsz=image_size, return_outputs=True) |
| object_prediction_list = [] |
| for _, image_results in enumerate(results): |
| if len(image_results)!=0: |
| image_predictions_in_xyxy_format = image_results['det'] |
| for pred in image_predictions_in_xyxy_format: |
| x1, y1, x2, y2 = ( |
| int(pred[0]), |
| int(pred[1]), |
| int(pred[2]), |
| int(pred[3]), |
| ) |
| bbox = [x1, y1, x2, y2] |
| score = pred[4] |
| category_name = model.model.names[int(pred[5])] |
| category_id = pred[5] |
| object_prediction = ObjectPrediction( |
| bbox=bbox, |
| category_id=int(category_id), |
| score=score, |
| category_name=category_name, |
| ) |
| object_prediction_list.append(object_prediction) |
|
|
| image = read_image(image) |
| output_image = visualize_object_predictions(image=image, object_prediction_list=object_prediction_list) |
| return output_image['image'] |
| |
|
|
| inputs = [ |
| gr.inputs.Image(type="filepath", label="Input Image"), |
| gr.inputs.Dropdown(["aijack/v8n", "aijack/v8m","aijack/v8x"], |
| default="aijack/v8m", label="Model"), |
| gr.inputs.Slider(minimum=320, maximum=1280, default=640, step=32, label="Image Size"), |
| gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.25, step=0.05, label="Confidence Threshold"), |
| gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.45, step=0.05, label="IOU Threshold"), |
| ] |
|
|
| outputs = gr.outputs.Image(type="filepath", label="Output Image") |
| title = "Ultralytics YOLOv8: State-of-the-Art YOLO Models" |
| article = "<p style='text-align: center'><a href='http://claireye.com.tw'>Claireye</a> | 2023</p>" |
| examples = [['test1.jpg', 'aijack/v8m', 640, 0.25, 0.45], ['test2.jpeg', 'aijack/v8x', 1280, 0.25, 0.45]] |
| demo_app = gr.Interface( |
| fn=yolov8_go, |
| inputs=inputs, |
| outputs=outputs, |
| title=title, |
| article=article, |
| examples=examples, |
| cache_examples=True |
| |
| ) |
| demo_app.launch(debug=True, enable_queue=True) |
