Update app.py

app.py

@@ -5,37 +5,25 @@ import sahi.predict
 import sahi.slicing
 from PIL import Image
 import numpy
+from huggingface_hub import hf_hub_download
+import torch
+
 
 IMAGE_SIZE = 640
 
-# 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",
-)
+model_path=hf_hub_download("kadirnar/deprem_model_v1", filename="last.pt",revision="main")
 
-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",
-)
 
+current_device='cuda' if torch.cuda.is_available() else 'cpu'
 
 # Model
 model = AutoDetectionModel.from_pretrained(
-    model_type="yolov5", model_path="yolov5s6.pt", device="cpu", confidence_threshold=0.5, image_size=IMAGE_SIZE
+    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,
@@ -47,34 +35,26 @@ def sahi_yolo_inference(
     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."
-        )
-
-    # 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_1 = Image.fromarray(visual_result_1["image"])
-
-    # sliced inference
-    prediction_result_2 = sahi.predict.get_sliced_prediction(
+    #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),
@@ -85,18 +65,33 @@ def sahi_yolo_inference(
         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,
-    )
+        )
+        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
 
-    output_2 = Image.fromarray(visual_result_2["image"])
 
-    return output_1, output_2
+    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"),
@@ -116,8 +111,7 @@ inputs = [
 ]
 
 outputs = [
-    gr.outputs.Image(type="pil", label="YOLOv5s"),
-    gr.outputs.Image(type="pil", label="YOLOv5s + SAHI"),
+    gr.outputs.Image(type="pil", label="Output")
 ]
 
 title = "Small Object Detection with SAHI + YOLOv5"