Update app.py

app.py

@@ -1,94 +1,63 @@
-import torch
 import gradio as gr
 import numpy as np
-from torchvision.ops import nms
-from PIL import Image
 import cv2
+import os
+from ultralytics import YOLO
+
+# Load the YOLO model
+model = YOLO('best.pt')
+
+# Function for image processing
+def show_preds_image(image_path):
+    image = cv2.imread(image_path)
+    results = model.predict(source=image_path)
+    annotated_image = results[0].plot()
+    return cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
+
+# Function for video processing
+def show_preds_video(video_path):
+    cap = cv2.VideoCapture(video_path)
+    out_frames = []
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+        results = model.predict(source=frame)
+        annotated_frame = results[0].plot()
+        out_frames.append(cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB))
+    cap.release()
+
+    # Save the annotated video
+    output_path = "annotated_video.mp4"
+    height, width, _ = out_frames[0].shape
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
+    for frame in out_frames:
+        writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+    writer.release()
+    return output_path
+
+# Gradio interfaces
+inputs_image = gr.Image(type="filepath", label="Input Image")
+outputs_image = gr.Image(type="numpy", label="Output Image")
+interface_image = gr.Interface(
+    fn=show_preds_image,
+    inputs=inputs_image,
+    outputs=outputs_image,
+    title="Taxi & License Plate Detection with Image"    
+)
 
-# Load the model
-model = torch.jit.load("best.torchscript")
-model.eval()
-
-# Define the detection function
-def detect_taxi_plate(image):
-    try:
-        # Preprocess the image
-        image_resized = Image.fromarray(image).resize((640, 640))
-        input_tensor = torch.from_numpy(np.array(image_resized).transpose(2, 0, 1) / 255.0).unsqueeze(0).float()
-
-        # Run inference
-        output = model(input_tensor)
-        detection_data = output[0][0].detach().numpy()  # Remove batch dimension
-         
-        # Filter detections by confidence threshold
-        filtered_detections = detection_data[detection_data[:, 4] >= 0.5]
-        # Define class names 
-        class_names = ["plate", "taxi"]
-
-        # Prepare boxes for NMS
-        boxes = []
-        confidences = []
-        labels = []
-        for detection in filtered_detections:
-            if len(detection) < 7:  # Ensure detection has enough elements
-                continue
-            x_center, y_center, width, height = detection[:4]
-            confidence = detection[4]
-            print(confidence)
-            class_probs = detection[5:]  # Probabilities for all classes
-
-            # Get the predicted class by finding the max probability index
-            class_index = np.argmax(class_probs)
-            class_label = class_names[class_index]
-            print(class_label)
-
-            x_min = int(x_center - width / 2.2)
-            y_min = int(y_center - height / 2.2)
-            x_max = int(x_center + width / 2.2)
-            y_max = int(y_center + height / 2.2)
-
-            boxes.append([x_min, y_min, x_max, y_max])
-            confidences.append(confidence)
-            labels.append(class_label)
-
-        if not boxes:  # No valid boxes
-            raise ValueError("No detections.")
-
-        boxes_tensor = torch.tensor(boxes, dtype=torch.float32)
-        scores_tensor = torch.tensor(confidences, dtype=torch.float32)
-
-        # Apply NMS
-        iou_threshold = 0.5
-        nms_indices = nms(boxes_tensor, scores_tensor, iou_threshold)
-        nms_boxes = boxes_tensor[nms_indices].tolist()
-        nms_labels = [labels[i] for i in nms_indices]
-
-        # Draw bounding boxes
-        image_with_boxes = image.copy()
-        for i, box in enumerate(nms_boxes):
-            x_min, y_min, x_max, y_max = map(int, box)
-            label = nms_labels[i]
-            cv2.rectangle(image_with_boxes, (x_min, y_min), (x_max, y_max), (255, 0, 0), 2)
-            cv2.putText(image_with_boxes, label, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
-
-        return image_with_boxes
-
-    except Exception as e:
-        print(f"Error: {str(e)}")
-        # Return error as text overlay on the image
-        image_with_error = image.copy()
-        cv2.putText(image_with_error, f"Error: {str(e)}", (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
-        return image_with_error
-
-# Define the Gradio interface
-interface = gr.Interface(
-    fn=detect_taxi_plate,
-    inputs=gr.Image(type="numpy", label="Upload Image"),
-    outputs=gr.Image(type="numpy", label="Output Image"),
-    title="ITI107 Assignment: Taxi & License Plate Detection",
-    description="Admin Number: 4744695Y\n\nUpload an image to detect if a Taxi and/or License Plate is present."
+inputs_video = gr.Video(label="Input Video")
+outputs_video = gr.Video(label="Annotated Output")
+interface_video = gr.Interface(
+    fn=show_preds_video,
+    inputs=inputs_video,
+    outputs=outputs_video,
+    title="Taxi & License Plate Detection with Video"
 )
 
-# Launch the app
-if __name__ == "__main__":
-    interface.launch(share=True)
+gr.TabbedInterface(
+    [interface_image, interface_video],
+    tab_names=['Image Inference', 'Video Inference']
+).launch(share=True)