Update app.py

app.py

@@ -1,62 +1,26 @@
 import gradio as gr
+import torch
 import cv2
-import time
 import numpy as np
-from sahi import AutoDetectionModel
-from sahi.predict import get_sliced_prediction
-from pathlib import Path
+from ultralytics import YOLO
 
-# Load the detection model
-detection_model = AutoDetectionModel.from_pretrained(
-    model_type='ultralytics',
-    model_path="./DDR.pt",  # Replace with your model path
-    confidence_threshold=0.01,
-    device="cpu"  # Change to 'cuda:0' if you have a GPU
-)
-
-OUTPUT_PATH = "./pred_image.jpg"
-TEMP_PNG_PATH = "./pred_image.png"
+def load_model(model_path="DDR.pt"):
+    return YOLO(model_path)
 
-def wait_for_file(file_path, timeout=10):
-    """Poll for the file to exist until the timeout (in seconds) is reached."""
-    start_time = time.time()
-    while not Path(file_path).exists():
-        if time.time() - start_time > timeout:
-            return False
-        time.sleep(0.5)
-    return True
+model = load_model()
 
-def run_inference(image):
-    # Perform sliced prediction on the input image.
-    result = get_sliced_prediction(
-        image,
-        detection_model,
-        slice_height=256,
-        slice_width=256,
-        overlap_height_ratio=0.2,
-        overlap_width_ratio=0.2
-    )
-    
-    # Export visualization to a temporary PNG file.
-    result.export_visuals(export_dir=Path(TEMP_PNG_PATH).parent, file_name=Path(TEMP_PNG_PATH).name)
-    
-    # Wait for the PNG file to be created.
-    if not wait_for_file(TEMP_PNG_PATH, timeout=10):
-        raise FileNotFoundError(f"SAHI did not save the PNG file at {TEMP_PNG_PATH}")
-    
-    # Read the PNG image, convert it to JPG, and remove the temporary file.
-    processed_image = cv2.imread(TEMP_PNG_PATH)
-    cv2.imwrite(OUTPUT_PATH, processed_image)
-    Path(TEMP_PNG_PATH).unlink()  # Delete the temporary PNG
-    
-    return OUTPUT_PATH
+def predict(image):
+    results = model(image, conf=0.01)
+    pred_img = results[0].plot()  # Visualize detections
+    return pred_img
 
-demo = gr.Interface(
-    fn=run_inference,
+iface = gr.Interface(
+    fn=predict,
     inputs=gr.Image(type="numpy"),
-    outputs=gr.Image(type="filepath"),
-    title="YOLO11 Object Detection",
-    description="Upload an image to run inference using YOLO11"
+    outputs=gr.Image(type="numpy"),
+    title="DDR-Detection",
+    description="Upload an image, and the model will detect objects using YOLO11.",
 )
 
-demo.launch(share=True)
+if __name__ == "__main__":
+    iface.launch()