Spaces:
Sleeping
Sleeping
File size: 5,024 Bytes
297a2c6 ce7d026 297a2c6 ce7d026 297a2c6 ce7d026 718b31c 297a2c6 ce7d026 297a2c6 ce7d026 297a2c6 ce7d026 718b31c ce7d026 297a2c6 ce7d026 718b31c ce7d026 718b31c ce7d026 718b31c ce7d026 718b31c ce7d026 297a2c6 718b31c 297a2c6 ce7d026 297a2c6 718b31c 297a2c6 718b31c 297a2c6 3e3c589 |
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 138 139 140 141 142 143 144 |
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_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.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."
# )
rect_th = None or max(round(sum(image.size) / 2 * 0.001), 1)
text_th = None or max(rect_th - 1, 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.inputs.Dropdown(choices=model_types,label="Choose Model Type",type="value",),
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 = [
[model_types[1],"satellite_original.tif", 256, 256, 0.2, 0.2, "GREEDYNMM", "IOS", 0.5, True],
[model_types[0],"26.jpg", 256, 256, 0.2, 0.2, "GREEDYNMM", "IOS", 0.5, True],
[model_types[0],"27.jpg", 512, 512, 0.2, 0.2, "GREEDYNMM", "IOS", 0.5, True],
[model_types[0],"28.jpg", 512, 512, 0.2, 0.2, "GREEDYNMM", "IOS", 0.5, True],
[model_types[0],"31.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)
|