Update app.py

app.py

@@ -3,24 +3,25 @@ from ultralytics import YOLO
 import cv2
 import tempfile
 
-# Function to load the model dynamically based on user input.
-def load_model(model_id: str):
+# Function to load a custom YOLO model from an uploaded file.
+def load_model(model_file):
     try:
-        model = YOLO(model_id)
+        # model_file is a TemporaryFile object. Use .name to get its path.
+        model = YOLO(model_file.name)
         return model
     except Exception as e:
         return f"Error loading model: {e}"
 
-# Inference for images: runs the model and plots predictions.
+# Function to perform inference on an image.
 def predict_image(model, image):
     try:
         results = model(image)
-        annotated_frame = results[0].plot()  # This handles detection, segmentation, or OBB models.
+        annotated_frame = results[0].plot()  # This should work across 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.
+# Function to perform inference on a video.
 def predict_video(model, video_file):
     try:
         cap = cv2.VideoCapture(video_file.name)
@@ -35,7 +36,7 @@ def predict_video(model, video_file):
         
         if not frames:
             return "Error: No frames processed from video."
-            
+        
         height, width, _ = frames[0].shape
         fourcc = cv2.VideoWriter_fourcc(*'mp4v')
         temp_video_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
@@ -47,10 +48,11 @@ 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):
-    model = load_model(model_id)
-    if isinstance(model, str):  # Indicates an error message.
+# Unified inference function that takes an uploaded model file, an input media file, and the selected media type.
+def inference(model_file, input_media, media_type):
+    model = load_model(model_file)
+    # Check if model loading resulted in an error message.
+    if isinstance(model, str):
         return model
 
     if media_type == "Image":
@@ -60,20 +62,24 @@ def inference(model_id, input_media, media_type):
     else:
         return "Unsupported 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.File(label="Upload Image/Video File")
+# Updated Gradio components:
+# - A file input for the custom YOLO model (.pt file)
+# - A file input for the image or video to process
+# - A radio button for selecting between image and video processing.
+model_file_input = gr.File(label="Upload Custom YOLO Model (.pt file)")
+media_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")
+output_component = gr.File(label="Processed Output")
 
-# Create Gradio interface.
+# Create the Gradio interface.
 iface = gr.Interface(
     fn=inference,
-    inputs=[model_text_input, file_input, media_type_dropdown],
-    outputs=output_file,
-    title="Dynamic Ultralytics YOLO Inference",
+    inputs=[model_file_input, media_file_input, media_type_dropdown],
+    outputs=output_component,
+    title="Custom YOLO Model Inference",
     description=(
-        "Enter the model identifier (supports detection, segmentation, or OBB models) and upload an image or video."
+        "Upload your custom YOLO model (for detection, segmentation, or OBB) along with an image or video file "
+        "to run inference. The system dynamically loads your model and processes the media accordingly."
     )
 )