Update app.py

app.py

@@ -15,7 +15,7 @@ BOX_PADDING = 2
 DETECTION_MODEL = YOLO("models/detector-model.pt") 
 CLASSIFICATION_MODEL = YOLO("models/classifier-model.pt") # 0: block, 1: checked, 2: unchecked
 
-def detect(image):
+def detect(image_path):
     """
     Output inference image with bounding box
 
@@ -24,6 +24,10 @@ def detect(image):
 
     Return: image with bounding boxes drawn 
     """
+    image = cv2.imread(image_path)
+    if image is None:
+        return image
+    
     # Predict on image
     results = DETECTION_MODEL.predict(source=image, conf=0.2, iou=0.8) # Predict on image
     boxes = results[0].boxes # Get bounding boxes
@@ -76,20 +80,4 @@ def detect(image):
 iface = gr.Interface(fn=detect, 
                      inputs=gr.inputs.Image(label="Upload scanned document", type="filepath"), 
                      outputs="image")
-iface.launch()
-
-
-
-
-
-
-
-
-            
-
-
-
-
-    
-
-
+iface.launch()