Update app.py

app.py

@@ -27,25 +27,30 @@ def detect_and_crop_document(image):
     results = modelY(image_np, conf=0.85, device='cpu')
     cropped_images = []
     predictions = []
+    
     for result in results:
         for box in result.boxes:
             x1, y1, x2, y2 = map(int, box.xyxy[0])
             conf = int(box.conf[0] * 100)  # Convert confidence to percentage
             cls = int(box.cls[0])
-            class_name = modelY.names[cls].capitalize()  # Ensure class names are capitalized
+            class_name = modelY.names[cls].capitalize()  # Capitalize class names
             cropped_image_np = image_np[y1:y2, x1:x2]
             cropped_image = Image.fromarray(cropped_image_np)
             cropped_images.append(cropped_image)
-            predictions.append(f"STNK {class_name} -- (Confidence: {conf}%)")
+            predictions.append(f"Detected: STNK {class_name} -- (Confidence: {conf}%)")
+    
     if not cropped_images:
         return None, "No document detected"
-    return cropped_images[0], predictions[0]
+    return cropped_images, predictions
 
 # Gradio interface
 def process_image(image):
     preprocessed_image = preprocessing(image)
-    cropped_image, prediction = detect_and_crop_document(preprocessed_image)
-    return cropped_image, prediction
+    cropped_images, predictions = detect_and_crop_document(preprocessed_image)
+    
+    if cropped_images:
+        return cropped_images, '\n'.join(predictions)
+    return None, "No document detected"
 
 with gr.Blocks(css=".gr-button {background-color: #4caf50; color: white; font-size: 16px; padding: 10px 20px; border-radius: 8px;}") as demo:
     gr.Markdown(
@@ -61,7 +66,7 @@ with gr.Blocks(css=".gr-button {background-color: #4caf50; color: white; font-si
                 clear_btn = gr.Button("Clear")
                 submit_btn = gr.Button("Detect Document")
         with gr.Column(scale=2):
-            output_image = gr.Image(label="Cropped Document", interactive=False)
+            output_image = gr.Gallery(label="Cropped Documents", interactive=False)
             output_text = gr.Textbox(label="Detection Result", interactive=False)
 
     submit_btn.click(process_image, inputs=input_image, outputs=[output_image, output_text])