File size: 4,603 Bytes
297a2c6
 
 
 
 
 
 
ce7d026
 
 
297a2c6
 
 
ce7d026
297a2c6
 
ce7d026
297a2c6
 
 
ce7d026
297a2c6
 
 
 
ce7d026
297a2c6
 
 
 
 
 
 
 
 
 
 
ce7d026
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
297a2c6
 
 
 
 
 
 
 
 
 
ce7d026
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
297a2c6
 
ce7d026
297a2c6
 
 
ce7d026
297a2c6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ce7d026
297a2c6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
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("kadirnar/deprem_model_v1", filename="last.pt",revision="main")


current_device='cuda' if torch.cuda.is_available() else 'cpu'

# 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.2,
    overlap_width_ratio=0.2,
    postprocess_type="GREEDYNMM",
    postprocess_match_metric="IOS",
    postprocess_match_threshold=0.5,
    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."
    #     )

    


    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,
        )
        output = Image.fromarray(visual_result_2["image"])

    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,
        )
        output = Image.fromarray(visual_result_1["image"])

    # sliced inference


    return output


inputs = [
    gr.Dropdown(choices=["YOLOv5","YOLOv5 + SAHI"],label="Choose Model Type"),
    gr.inputs.Image(type="pil", label="Original Image"),
    gr.inputs.Number(default=512, label="slice_height"),
    gr.inputs.Number(default=512, label="slice_width"),
    gr.inputs.Number(default=0.2, label="overlap_height_ratio"),
    gr.inputs.Number(default=0.2, label="overlap_width_ratio"),
    gr.inputs.Dropdown(
        ["NMS", "GREEDYNMM"],
        type="value",
        default="GREEDYNMM",
        label="postprocess_type",
    ),
    gr.inputs.Dropdown(
        ["IOU", "IOS"], type="value", default="IOS", label="postprocess_type"
    ),
    gr.inputs.Number(default=0.5, label="postprocess_match_threshold"),
    gr.inputs.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 = [
    ["apple_tree.jpg", 256, 256, 0.2, 0.2, "GREEDYNMM", "IOS", 0.5, True],
    ["highway.jpg", 256, 256, 0.2, 0.2, "GREEDYNMM", "IOS", 0.5, True],
    ["highway2.jpg", 512, 512, 0.2, 0.2, "GREEDYNMM", "IOS", 0.5, True],
    ["highway3.jpg", 512, 512, 0.2, 0.2, "GREEDYNMM", "IOS", 0.5, True],
]

gr.Interface(
    sahi_yolo_inference,
    inputs,
    outputs,
    title=title,
    description=description,
    article=article,
    examples=examples,
    theme="huggingface",
).launch(debug=True, enable_queue=True)