Use separate tab interface for image & video UI

app.py

@@ -1,68 +1,88 @@
-from ultralytics import YOLO
 import gradio as gr
+from ultralytics import YOLO
+from PIL import Image
+import os
 import cv2
 import tempfile
-import os
-import subprocess
 
 # Load YOLOv8 model
-model_path = "./best_yolo8_model/best.pt"  # Adjust to your trained YOLOv8 model path
-detection_model = YOLO(model_path)
+model = YOLO('./best_yolo8_model/best.pt')
+
+# Define fixed dimensions
+image_height = 300
+image_width = 400
+video_height = 300
+video_width = 400
+
+def process_image(image):
+    # Perform inference on the uploaded image
+    results = model.predict(source=image, conf=0.5)
+    result_img = results[0].plot()
+    # Ensure image color stays consistent
+    return Image.fromarray(cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB))
 
-# Function to process videos
 def process_video(video_path):
-    # Read the video
+    # Perform inference on the uploaded video
+    temp_dir = tempfile.mkdtemp()
+    output_video_path = os.path.join(temp_dir, "processed_video.mp4")
     cap = cv2.VideoCapture(video_path)
-    if not cap.isOpened():
-        return "Error: Cannot open video file."
-
-    # Get video properties
-    fps = int(cap.get(cv2.CAP_PROP_FPS))
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))
     
-    # Define a temporary raw output filename
-    raw_output = tempfile.NamedTemporaryFile(suffix=".avi", delete=False).name
-
-    # Define codec and output format for intermediate video
-    fourcc = cv2.VideoWriter_fourcc(*'XVID')  # Intermediate codec
-    out = cv2.VideoWriter(raw_output, fourcc, fps, (width, height))
-
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
             break
-        
-        # Run YOLO inference
-        results = detection_model.predict(frame, conf=0.5, iou=0.6)
-        
-        # Draw bounding boxes
-        annotated_frame = results[0].plot()
-
-        # Write the annotated frame to the intermediate output video
-        out.write(annotated_frame)
+        results = model.predict(source=frame, conf=0.5)
+        result_frame = results[0].plot()
+        # Ensure color correction for the video frame
+        result_frame_corrected = cv2.cvtColor(result_frame, cv2.COLOR_BGR2RGB)
+        out.write(cv2.cvtColor(result_frame_corrected, cv2.COLOR_RGB2BGR))  # Maintain video color format
     
     cap.release()
     out.release()
+    return output_video_path
 
-    # Convert the intermediate video to MP4 with H.264 codec for browser playback
-    final_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
-    command = [
-        "ffmpeg", "-i", raw_output, "-vcodec", "libx264", "-preset", "fast", "-crf", "23", "-y", final_output
-    ]
-    subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-
-    # Clean up raw output
-    os.remove(raw_output)
-
-    return final_output
-
-# Gradio interface for video processing
-iface = gr.Interface(
-    fn=process_video,
-    inputs=gr.Video(label="Input Video"),  # Fixed input video component
-    outputs=gr.Video(label="Processed Video"),  # Fixed output video component
-    title="YOLOv8 Object Detection - Video Processing"
-)
+# Create Gradio interface with tabs
+with gr.Blocks() as app:
+    gr.Markdown("## YOLOv8 Object Detection - Image & Video")
+    gr.Markdown(
+        "This app detects objects in images and videos using a YOLOv8s model. It detects DURIAN and RAMBUTAN fruits. Use the tabs below to process images or videos."
+    )
+    
+    with gr.Tabs():
+        with gr.TabItem("Image Detection"):
+            with gr.Row():
+                with gr.Column():
+                    image_input = gr.Image(
+                        type="pil", label="Input Image", elem_id="image_input", 
+                        width=image_width, height=image_height
+                    )
+                with gr.Column():
+                    image_output = gr.Image(
+                        type="pil", label="Output Image", elem_id="image_output", 
+                        width=image_width, height=image_height
+                    )
+            image_submit = gr.Button("Detect Objects in Image")
+            image_submit.click(process_image, inputs=image_input, outputs=image_output)
+        
+        with gr.TabItem("Video Detection"):
+            with gr.Row():
+                with gr.Column():
+                    video_input = gr.Video(
+                        label="Input Video", elem_id="video_input", 
+                        width=video_width, height=video_height
+                    )
+                with gr.Column():
+                    video_output = gr.Video(
+                        label="Output Video", elem_id="video_output", 
+                        width=video_width, height=video_height
+                    )
+            video_submit = gr.Button("Detect Objects in Video")
+            video_submit.click(process_video, inputs=video_input, outputs=video_output)
 
-iface.launch(share=True)
+app.launch()

requirements.txt

@@ -2,4 +2,4 @@ ultralytics
 gradio
 huggingface_hub
 pillow
-openvino
+ffmpeg-python