Update app.py

app.py

@@ -3,25 +3,28 @@ from ultralytics import YOLO
 import cv2
 import tempfile
 
-# Function to load a custom YOLO model from an uploaded file.
+# Load a custom YOLO model from the uploaded file.
 def load_model(model_file):
     try:
-        # 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}"
 
-# Function to perform inference on an image.
+# Run inference on an image and write the output to a PNG file.
 def predict_image(model, image):
     try:
         results = model(image)
-        annotated_frame = results[0].plot()  # This should work across detection, segmentation, or OBB models.
-        return annotated_frame
+        annotated_frame = results[0].plot()  # Works for detection, segmentation, and OBB models.
+        # Write annotated image to a temporary file.
+        tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
+        cv2.imwrite(tmp.name, annotated_frame)
+        return tmp.name
     except Exception as e:
         return f"Error during image inference: {e}"
 
-# Function to perform inference on a video.
+# Run inference on a video by processing frame-by-frame,
+# and write the annotated video to an MP4 file.
 def predict_video(model, video_file):
     try:
         cap = cv2.VideoCapture(video_file.name)
@@ -33,53 +36,55 @@ def predict_video(model, video_file):
             frames.append(annotated_frame)
             success, frame = cap.read()
         cap.release()
-        
         if not frames:
-            return "Error: No frames processed from video."
-        
+            return f"Error: No frames processed"
         height, width, _ = frames[0].shape
-        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-        temp_video_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
-        out = cv2.VideoWriter(temp_video_file.name, fourcc, 20.0, (width, height))
+        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+        tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
+        out = cv2.VideoWriter(tmp.name, fourcc, 20.0, (width, height))
         for frame in frames:
             out.write(frame)
         out.release()
-        return temp_video_file.name
+        return tmp.name
     except Exception as e:
         return f"Error during video inference: {e}"
 
-# Unified inference function that takes an uploaded model file, an input media file, and the selected media type.
+# Main inference function: loads the custom model and processes the input media.
+# Returns a tuple: (annotated_image, annotated_video).
+# One element will be a file path and the other None, based on the 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
-
+        # An error occurred during model loading.
+        return (model, None)
+    
     if media_type == "Image":
-        return predict_image(model, input_media)
+        out_image = predict_image(model, input_media)
+        return (out_image, None)
     elif media_type == "Video":
-        return predict_video(model, input_media)
+        out_video = predict_video(model, input_media)
+        return (None, out_video)
     else:
-        return "Unsupported media type."
+        return ("Unsupported media type", None)
 
-# 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.
+# Define Gradio interface components.
 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_component = gr.File(label="Processed Output")
 
-# Create the Gradio interface.
+# Define two outputs: one for images and one for videos.
+output_image = gr.Image(label="Annotated Image")
+output_video = gr.Video(label="Annotated Video")
+
+# Create a Gradio interface that returns a tuple: (image, video).
 iface = gr.Interface(
     fn=inference,
     inputs=[model_file_input, media_file_input, media_type_dropdown],
-    outputs=output_component,
+    outputs=[output_image, output_video],
     title="Custom YOLO Model Inference",
     description=(
-        "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."
+        "Upload your custom YOLO model (detection, segmentation, or OBB) along with an image or video file "
+        "to run inference. The system loads your model dynamically, processes the media, and displays the output."
     )
 )