Spaces:
Running
on
T4
Running
on
T4
| import gradio as gr | |
| import sahi.utils | |
| from sahi import AutoDetectionModel | |
| import sahi.predict | |
| import sahi.slicing | |
| from PIL import Image | |
| import numpy | |
| from huggingface_hub import hf_hub_download | |
| import torch | |
| IMAGE_SIZE = 640 | |
| model_path = hf_hub_download( | |
| "deprem-ml/Binafarktespit-yolo5x-v1-xview", filename="last.pt", revision="main" | |
| ) | |
| current_device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model_types = ["YOLOv5", "YOLOv5 + SAHI"] | |
| # Model | |
| model = AutoDetectionModel.from_pretrained( | |
| model_type="yolov5", | |
| model_path=model_path, | |
| device=current_device, | |
| confidence_threshold=0.5, | |
| image_size=IMAGE_SIZE, | |
| ) | |
| def sahi_yolo_inference( | |
| model_type, | |
| image, | |
| 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, | |
| ): | |
| # image_width, image_height = image.size | |
| # sliced_bboxes = sahi.slicing.get_slice_bboxes( | |
| # image_height, | |
| # image_width, | |
| # slice_height, | |
| # slice_width, | |
| # False, | |
| # overlap_height_ratio, | |
| # overlap_width_ratio, | |
| # ) | |
| # if len(sliced_bboxes) > 60: | |
| # raise ValueError( | |
| # f"{len(sliced_bboxes)} slices are too much for huggingface spaces, try smaller slice size." | |
| # ) | |
| rect_th = None or max(round(sum(image.size) / 2 * 0.0001), 1) | |
| text_th = None or max(rect_th - 2, 1) | |
| if "SAHI" in model_type: | |
| prediction_result_2 = sahi.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 = sahi.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 | |
| else: | |
| # standard inference | |
| prediction_result_1 = sahi.predict.get_prediction( | |
| image=image, detection_model=model | |
| ) | |
| print(image) | |
| visual_result_1 = sahi.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 | |
| # sliced inference | |
| inputs = [ | |
| gr.Dropdown( | |
| choices=model_types, | |
| label="Choose Model Type", | |
| type="value", | |
| value=model_types[1], | |
| ), | |
| gr.Image(type="pil", label="Original Image"), | |
| gr.Number(default=512, label="slice_height"), | |
| gr.Number(default=512, label="slice_width"), | |
| gr.Number(default=0.1, label="overlap_height_ratio"), | |
| gr.Number(default=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.Number(default=0.5, label="postprocess_match_threshold"), | |
| gr.Checkbox(default=True, label="postprocess_class_agnostic"), | |
| ] | |
| outputs = [gr.outputs.Image(type="pil", label="Output")] | |
| title = "Small Object Detection with SAHI + YOLOv5" | |
| description = "SAHI + YOLOv5 demo for small object detection. 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 = [ | |
| [model_types[1], "26.jpg", 512, 512, 0.1, 0.1, "NMS", "IOU", 0.25, False], | |
| [model_types[1], "27.jpg", 512, 512, 0.1, 0.1, "NMS", "IOU", 0.25, False], | |
| [model_types[1], "28.jpg", 512, 512, 0.1, 0.1, "NMS", "IOU", 0.25, False], | |
| [model_types[1], "31.jpg", 512, 512, 0.1, 0.1, "NMS", "IOU", 0.25, False], | |
| ] | |
| gr.Interface( | |
| sahi_yolo_inference, | |
| inputs, | |
| outputs, | |
| title=title, | |
| description=description, | |
| article=article, | |
| examples=examples, | |
| theme="huggingface", | |
| ).launch(debug=True, enable_queue=True) | |