Update app.py

app.py

@@ -6,6 +6,7 @@ from sahi import AutoDetectionModel
 from sahi.predict import get_sliced_prediction
 from pathlib import Path
 
+# Load the detection model
 detection_model = AutoDetectionModel.from_pretrained(
     model_type='ultralytics',
     model_path="./DDR.pt",  # Replace with your model path
@@ -16,7 +17,17 @@ detection_model = AutoDetectionModel.from_pretrained(
 OUTPUT_PATH = "./pred_image.jpg"
 TEMP_PNG_PATH = "./pred_image.png"
 
+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
+
 def run_inference(image):
+    # Perform sliced prediction on the input image.
     result = get_sliced_prediction(
         image,
         detection_model,
@@ -26,16 +37,17 @@ def run_inference(image):
         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)
     
-    time.sleep(2)
-    
-    if not Path(TEMP_PNG_PATH).exists():
+    # 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()
+    Path(TEMP_PNG_PATH).unlink()  # Delete the temporary PNG
     
     return OUTPUT_PATH
 
@@ -44,7 +56,7 @@ demo = gr.Interface(
     inputs=gr.Image(type="numpy"),
     outputs=gr.Image(type="filepath"),
     title="YOLO11 Object Detection",
-    description="Upload a DDR image to run inference using YOLO11"
+    description="Upload an image to run inference using YOLO11"
 )
 
-demo.launch()
+demo.launch(share=True)