File size: 4,505 Bytes
2f356cf
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import gradio as gr
import yolov5
import sahi.utils
import sahi.model
import sahi.predict
from PIL import Image
import numpy

# Images
sahi.utils.file.download_from_url(
    "https://user-images.githubusercontent.com/34196005/142730935-2ace3999-a47b-49bb-83e0-2bdd509f1c90.jpg",
    "apple_tree.jpg",
)
sahi.utils.file.download_from_url(
    "https://user-images.githubusercontent.com/34196005/142730936-1b397756-52e5-43be-a949-42ec0134d5d8.jpg",
    "highway.jpg",
)

sahi.utils.file.download_from_url(
    "https://user-images.githubusercontent.com/34196005/142742871-bf485f84-0355-43a3-be86-96b44e63c3a2.jpg",
    "highway2.jpg",
)

sahi.utils.file.download_from_url(
    "https://user-images.githubusercontent.com/34196005/142742872-1fefcc4d-d7e6-4c43-bbb7-6b5982f7e4ba.jpg",
    "highway3.jpg",
)


# Model
model = sahi.model.Yolov5DetectionModel(
    model_path="yolov5s6.pt", device="cpu", confidence_threshold=0.5
)


def sahi_yolo_inference(
    image,
    slice_height=512,
    slice_width=512,
    overlap_height_ratio=0.2,
    overlap_width_ratio=0.2,
    image_size=640,
    postprocess_type="UNIONMERGE",
    postprocess_match_metric="IOS",
    postprocess_match_threshold=0.5,
    postprocess_class_agnostic=False,
):

    # standard inference
    prediction_result_1 = sahi.predict.get_prediction(
        image=image, detection_model=model, image_size=image_size
    )
    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_1 = Image.fromarray(visual_result_1["image"])

    # sliced inference
    prediction_result_2 = sahi.predict.get_sliced_prediction(
        image=image,
        detection_model=model,
        image_size=image_size,
        slice_height=slice_height,
        slice_width=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_2 = Image.fromarray(visual_result_2["image"])

    return output_1, output_2


inputs = [
    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.Number(default=640, label="image_size"),
    gr.inputs.Dropdown(
        ["NMS", "UNIONMERGE"],
        type="value",
        default="UNIONMERGE",
        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="Standard YOLOv5s Inference"),
    gr.outputs.Image(type="pil", label="Sliced YOLOv5s Inference"),
]

title = "SAHI + YOLOv5"
description = "SAHI + YOLOv5 Gradio demo for 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, 640, "UNIONMERGE", "IOS", 0.5, True],
    ["highway.jpg", 256, 256, 0.2, 0.2, 640, "UNIONMERGE", "IOS", 0.5, True],
    ["highway2.jpg", 512, 512, 0.2, 0.2, 640, "UNIONMERGE", "IOS", 0.5, True],
    ["highway3.jpg", 1024, 1024, 0.2, 0.2, 640, "UNIONMERGE", "IOS", 0.5, True],
]

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