Update app.py

app.py

@@ -15,7 +15,7 @@ model_path=hf_hub_download("kadirnar/deprem_model_v1", filename="last.pt",revisi
 
 
 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
@@ -50,8 +50,9 @@ def sahi_yolo_inference(
     #         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(
@@ -69,8 +70,11 @@ def sahi_yolo_inference(
         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
@@ -81,17 +85,19 @@ def sahi_yolo_inference(
         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
 
 
-    return output
 
 
 inputs = [
-    gr.Dropdown(choices=["YOLOv5","YOLOv5 + SAHI"],label="Choose Model Type"),
+    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"),
@@ -118,12 +124,13 @@ 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],
-]
+    [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,
@@ -133,4 +140,4 @@ gr.Interface(
     article=article,
     examples=examples,
     theme="huggingface",
-).launch(debug=True, enable_queue=True)
+).launch(debug=True, enable_queue=True,live=False)