Update app.py

app.py

@@ -28,7 +28,6 @@ category_dict = {
     77: 'teddy bear', 78: 'hair drier', 79: 'toothbrush'
 }
 
-
 @spaces.GPU(duration=200)
 def yolov10_inference(image, model_id, image_size, conf_threshold, iou_threshold):
     model_path = download_models(model_id)
@@ -44,6 +43,11 @@ def yolov10_inference(image, model_id, image_size, conf_threshold, iou_threshold
 
     return annotated_image
 
+def yolov10_inference_multi(image, image_size, conf_threshold, iou_threshold):
+    yolov10s_image = yolov10_inference(image, "yolov10s.pt", image_size, conf_threshold, iou_threshold)
+    yolov10m_image = yolov10_inference(image, "yolov10m.pt", image_size, conf_threshold, iou_threshold)
+    return yolov10s_image, yolov10m_image
+
 def app():
     with gr.Blocks():
         with gr.Row():
@@ -87,6 +91,8 @@ def app():
 
             with gr.Column():
                 output_image = gr.Image(type="pil", label="Annotated Image")
+                output_image_s = gr.Image(type="pil", label="Annotated Image yolov10s")
+                output_image_m = gr.Image(type="pil", label="Annotated Image yolov10m")
 
         yolov10_infer.click(
             fn=yolov10_inference,
@@ -100,6 +106,17 @@ def app():
             outputs=[output_image],
         )
 
+        yolov10_infer.click(
+            fn=yolov10_inference_multi,
+            inputs=[
+                image,
+                image_size,
+                conf_threshold,
+                iou_threshold,
+            ],
+            outputs=[output_image_s, output_image_m],
+        )
+
         gr.Examples(
             examples=[
                 [
@@ -148,4 +165,4 @@ with gradio_app:
         with gr.Column():
             app()
 
-gradio_app.launch(debug=True)
+gradio_app.launch(debug=True)