Spaces:
Sleeping
Sleeping
| from sahi import utils, predict, AutoDetectionModel | |
| from PIL import Image | |
| import gradio as gr | |
| import numpy | |
| import torch | |
| import os | |
| os.system('pip install git+https://github.com/fcakyon/ultralyticsplus.git') | |
| model_id_list = ['deprem-ml/Binafarktespit-yolo5x-v1-xview', 'SerdarHelli/deprem_satellite_labeled_yolov8'] | |
| current_device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model_types = ["YOLOv5", "YOLOv5 + SAHI", "YOLOv8"] | |
| def sahi_yolov5_inference( | |
| image, | |
| model_id, | |
| model_type, | |
| image_size, | |
| slice_height=512, | |
| slice_width=512, | |
| overlap_height_ratio=0.1, | |
| overlap_width_ratio=0.1, | |
| postprocess_type="NMS", | |
| postprocess_match_metric="IOU", | |
| postprocess_match_threshold=0.25, | |
| postprocess_class_agnostic=False, | |
| ): | |
| rect_th = None or max(round(sum(image.size) / 2 * 0.0001), 1) | |
| text_th = None or max(rect_th - 2, 1) | |
| if model_type == "YOLOv5": | |
| # standard inference | |
| model = AutoDetectionModel.from_pretrained( | |
| model_type="yolov5", | |
| model_path=model_id, | |
| device=current_device, | |
| confidence_threshold=0.5, | |
| image_size=image_size, | |
| ) | |
| prediction_result_1 = predict.get_prediction( | |
| image=image, detection_model=model | |
| ) | |
| visual_result_1 = utils.cv.visualize_object_predictions( | |
| image=numpy.array(image), | |
| object_prediction_list=prediction_result_1.object_prediction_list, | |
| rect_th=rect_th, | |
| text_th=text_th, | |
| ) | |
| output = Image.fromarray(visual_result_1["image"]) | |
| return output | |
| elif model_type == "YOLOv5 + SAHI": | |
| model = AutoDetectionModel.from_pretrained( | |
| model_type="yolov5", | |
| model_path=model_id, | |
| device=current_device, | |
| confidence_threshold=0.5, | |
| image_size=image_size, | |
| ) | |
| prediction_result_2 = predict.get_sliced_prediction( | |
| image=image, | |
| detection_model=model, | |
| slice_height=int(slice_height), | |
| slice_width=int(slice_width), | |
| overlap_height_ratio=overlap_height_ratio, | |
| overlap_width_ratio=overlap_width_ratio, | |
| postprocess_type=postprocess_type, | |
| postprocess_match_metric=postprocess_match_metric, | |
| postprocess_match_threshold=postprocess_match_threshold, | |
| postprocess_class_agnostic=postprocess_class_agnostic, | |
| ) | |
| visual_result_2 = utils.cv.visualize_object_predictions( | |
| image=numpy.array(image), | |
| object_prediction_list=prediction_result_2.object_prediction_list, | |
| rect_th=rect_th, | |
| text_th=text_th, | |
| ) | |
| output = Image.fromarray(visual_result_2["image"]) | |
| return output | |
| elif model_type == "YOLOv8": | |
| from ultralyticsplus import YOLO, render_result | |
| model = YOLO('SerdarHelli/deprem_satellite_labeled_yolov8') | |
| result = model.predict(image, imgsz=image_size)[0] | |
| render = render_result(model=model, image=image, result=result, rect_th=rect_th, text_th=text_th) | |
| return render | |
| inputs = [ | |
| gr.Image(type="pil", label="Original Image"), | |
| gr.Dropdown(choices=model_id_list,label="Choose Model",value=model_id_list[0]), | |
| gr.Dropdown( choices=model_types, label="Choose Model Type", value=model_types[1]), | |
| gr.Slider(minimum=128, maximum=2048, value=640, step=32, label="Image Size"), | |
| gr.Slider(minimum=128, maximum=2048, value=512, step=32, label="Slice Height"), | |
| gr.Slider(minimum=128, maximum=2048, value=512, step=32, label="Slice Width"), | |
| gr.Slider(minimum=0.0, maximum=1.0, value=0.1, step=0.1, label="Overlap Height Ratio"), | |
| gr.Slider(minimum=0.0, maximum=1.0, value=0.1, step=0.1, label="Overlap Width Ratio"), | |
| gr.Dropdown(["NMS", "GREEDYNMM"], type="value", value="NMS", label="Postprocess Type"), | |
| gr.Dropdown(["IOU", "IOS"], type="value", value="IOU", label="Postprocess Type"), | |
| gr.Slider(minimum=0.0, maximum=1.0, value=0.5, step=0.1, label="Postprocess Match Threshold"), | |
| gr.Checkbox(value=True, label="Postprocess Class Agnostic"), | |
| ] | |
| outputs = [gr.outputs.Image(type="pil", label="Output")] | |
| title = "Building Detection from Satellite Images with SAHI + YOLOv5" | |
| description = "SAHI + YOLOv5 demo for building detection from satellite images. Upload an image or click an example image to use." | |
| article = "<p style='text-align: center'>SAHI is a lightweight vision library for performing large scale object detection/ instance segmentation.. <a href='https://github.com/obss/sahi'>SAHI Github</a> | <a href='https://medium.com/codable/sahi-a-vision-library-for-performing-sliced-inference-on-large-images-small-objects-c8b086af3b80'>SAHI Blog</a> | <a href='https://github.com/fcakyon/yolov5-pip'>YOLOv5 Github</a> </p>" | |
| examples = [ | |
| ["data/26.jpg", 'deprem-ml/Binafarktespit-yolo5x-v1-xview', "YOLOv5 + SAHI", 640, 512, 512, 0.1, 0.1, "NMS", "IOU", 0.25, False], | |
| ["data/27.jpg", 'deprem-ml/Binafarktespit-yolo5x-v1-xview', "YOLOv5 + SAHI", 640, 512, 512, 0.1, 0.1, "NMS", "IOU", 0.25, False], | |
| ["data/28.jpg", 'deprem-ml/Binafarktespit-yolo5x-v1-xview', "YOLOv5 + SAHI", 640, 512, 512, 0.1, 0.1, "NMS", "IOU", 0.25, False], | |
| ["data/31.jpg", 'deprem-ml/SerdarHelli-yolov8-v1-xview', "YOLOv8", 640, 512, 512, 0.1, 0.1, "NMS", "IOU", 0.25, False], | |
| ] | |
| demo = gr.Interface( | |
| sahi_yolov5_inference, | |
| inputs, | |
| outputs, | |
| title=title, | |
| description=description, | |
| article=article, | |
| examples=examples, | |
| theme="huggingface", | |
| cache_examples=True, | |
| ) | |
| demo.launch(debug=True, enable_queue=True) | |