| | 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_inference( |
| | 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.25, |
| | ): |
| | """ |
| | YOLOv8 inference function |
| | Args: |
| | image: Input image |
| | model_path: Path to the model |
| | image_size: Image size |
| | conf_threshold: Confidence threshold |
| | iou_threshold: IOU threshold |
| | Returns: |
| | Rendered image |
| | """ |
| | 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): |
| | 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(["Asma/GreenHawk_test"], |
| | default="Asma/GreenHawk_test", 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.25, step=0.05, label="IOU Threshold"), |
| | ] |
| |
|
| | outputs = gr.outputs.Image(type="filepath", label="Output Image") |
| | title = "GreenHawk - Visual Pollution Detection" |
| |
|
| | demo_app = gr.Interface( |
| | fn=yolov8_inference, |
| | inputs=inputs, |
| | outputs=outputs, |
| | title=title, |
| | |
| | cache_examples=True, |
| | theme='default', |
| | ) |
| | demo_app.launch(debug=True, enable_queue=True) |
