Update app.py

app.py

@@ -55,36 +55,32 @@ for model in models:
       model["cfg"].MODEL.DEVICE = "cpu"
 
 
-def inference(image, min_score, model_name):
-    results = []
-    
+def inference(image_url, image, min_score, model_name):
+    if image_url:
+        r = requests.get(image_url)
+        if r:
+            im = np.frombuffer(r.content, dtype="uint8")
+            im = cv2.imdecode(im, cv2.IMREAD_COLOR_BGR2RGB)
+    else:
+        # Model expect BGR!
+        im = image[:,:,::-1]
+
     model_id = model_name_to_id[model_name]
+
     models[model_id]["cfg"].MODEL.ROI_HEADS.SCORE_THRESH_TEST = min_score
     predictor = DefaultPredictor(models[model_id]["cfg"])
 
-    # Asegurémonos de que image es una lista de imágenes o solo una
-    if isinstance(image, list):  # Si es una lista de imágenes
-        images = image
-    else:  # Si es solo una imagen, la convertimos en una lista
-        images = [image]
-
-    for img in images:
-        # Convert image from BGR to RGB if necessary
-        im = img[:,:,::-1]
+    outputs = predictor(im)
 
-        # Make the prediction
-        outputs = predictor(im)
+    v = Visualizer(im, models[model_id]["metadata"], scale=1.2)
+    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
+        
+    # Convertir la imagen de salida de BGR a RGB
+    result_image = out.get_image()  # Esto sigue estando en BGR
+    result_image_rgb = result_image[:, :, ::-1]  # Convertir BGR a RGB
 
-        v = Visualizer(im, models[model_id]["metadata"], scale=1.2)
-        out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
+    return result_image_rgb
 
-        # Convert the result from BGR to RGB
-        result_image = out.get_image()
-        result_image_rgb = result_image[:, :, ::-1]  # Convert BGR to RGB
-
-        results.append(result_image_rgb)  # Add the processed image to the list
-
-    return results  # Return all the results
 
 title = "# DBMDZ Detectron2 Model Demo"
 description = """
@@ -99,18 +95,21 @@ with gr.Blocks() as demo:
     gr.Markdown(title)
     gr.Markdown(description)
     
+    with gr.Tab("From URL"):
+        url_input = gr.Textbox(label="Image URL", placeholder="https://api.digitale-sammlungen.de/iiif/image/v2/bsb10483966_00008/full/500,/0/default.jpg")
+    
     with gr.Tab("From Image"):
-        image_input = gr.Gallery(type="numpy", label="Input Images", elem_id="input_image")
+        image_input = gr.Image(type="numpy", label="Input Image")
 
     min_score = gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="Minimum score")
 
     model_name = gr.Radio(choices=[model["name"] for model in models], value=models[0]["name"], label="Select Detectron2 model")
 
-    output_gallery = gr.Gallery(type ="numpy", label="Output Images", elem_id="output_images")
+    output_image = gr.Image(type="pil", label="Output")
 
     inference_button = gr.Button("Submit")
     
-    inference_button.click(fn=inference, inputs=[image_input, min_score, model_name], outputs=output_gallery)
+    inference_button.click(fn=inference, inputs=[url_input, image_input, min_score, model_name], outputs=output_image)
 
     gr.Markdown(footer)