Update app.py

app.py

@@ -3,7 +3,7 @@ from ultralytics import YOLO
 import cv2
 import tempfile
 
-# Dynamic model loader: given a model identifier or path, it returns a YOLO model.
+# Function to load the model dynamically based on user input.
 def load_model(model_id: str):
     try:
         model = YOLO(model_id)
@@ -11,16 +11,16 @@ def load_model(model_id: str):
     except Exception as e:
         return f"Error loading model: {e}"
 
+# Inference for images: runs the model and plots predictions.
 def predict_image(model, image):
     try:
-        # Run inference on the image
         results = model(image)
-        # The plot() method overlays predictions (works across detection, segmentation, and OBB models)
-        annotated_frame = results[0].plot()
+        annotated_frame = results[0].plot()  # This handles detection, segmentation, or OBB models.
         return annotated_frame
     except Exception as e:
         return f"Error during image inference: {e}"
 
+# Inference for videos: processes the video frame by frame.
 def predict_video(model, video_file):
     try:
         cap = cv2.VideoCapture(video_file.name)
@@ -47,14 +47,12 @@ def predict_video(model, video_file):
     except Exception as e:
         return f"Error during video inference: {e}"
 
+# Unified inference function based on media type.
 def inference(model_id, input_media, media_type):
-    # Dynamically load the model using the provided identifier
     model = load_model(model_id)
-    if isinstance(model, str):
-        # This means an error message was returned
+    if isinstance(model, str):  # Indicates an error message.
         return model
 
-    # Process based on media type
     if media_type == "Image":
         return predict_image(model, input_media)
     elif media_type == "Video":
@@ -62,24 +60,20 @@ def inference(model_id, input_media, media_type):
     else:
         return "Unsupported media type."
 
-# Create Gradio inputs:
-# - A text input for the model identifier
-# - A file input for images or videos
-# - A radio selection for choosing the media type
+# Updated Gradio interface components using the new API.
 model_text_input = gr.Textbox(label="Model Identifier", placeholder="e.g., yolov8n.pt, yolov8-seg.pt, yolov8-obb.pt")
-file_input = gr.inputs.File(label="Upload Image/Video File")
+file_input = gr.File(label="Upload Image/Video File")
 media_type_dropdown = gr.Radio(choices=["Image", "Video"], label="Select Media Type", value="Image")
+output_file = gr.File(label="Processed Output")
 
-# Decide on output type: for images, use an image output; for video, use a file output.
-# For a combined interface, we use File output because it can handle both cases with minimal changes.
+# Create Gradio interface.
 iface = gr.Interface(
     fn=inference,
     inputs=[model_text_input, file_input, media_type_dropdown],
-    outputs=gr.outputs.File(label="Processed Output"),
+    outputs=output_file,
     title="Dynamic Ultralytics YOLO Inference",
     description=(
-        "Enter the model identifier (for detection, segmentation, or OBB models) and upload an image or video "
-        "to run inference. The system dynamically loads the specified model and processes your media input."
+        "Enter the model identifier (supports detection, segmentation, or OBB models) and upload an image or video."
     )
 )