Upload app.py

app.py

@@ -0,0 +1,220 @@
+# -*- coding: utf-8 -*-
+"""Judol Gradio YOLO11.ipynb
+
+Automatically generated by Colab.
+
+Original file is located at
+    https://colab.research.google.com/drive/1oiuTAi-cys1ydtUhSDJSRdeA02mAmZQH
+"""
+
+
+
+import cv2
+from ultralytics import YOLO
+import gradio as gr
+import imageio
+
+
+
+
+model = YOLO('https://huggingface.co/JrEasy/Judol-Detection-YOLO11/resolve/main/best.pt')
+
+
+confidence_threshold = 0.6
+
+class_names = {
+    0: "BK8",
+    1: "Gate of Olympus",
+    2: "Princess",
+    3: "Starlight Princess",
+    4: "Zeus",
+}
+
+class_colors = {
+    0: (0, 255, 0),       # Green for BK8
+    1: (255, 0, 0),       # Blue for Gate of Olympus
+    2: (0, 0, 255),       # Red for Princess
+    3: (255, 255, 0),     # Cyan for Starlight Princess
+    4: (255, 0, 255),     # Magenta for Zeus
+}
+
+def format_time_ranges(timestamps, classes):
+
+    if not timestamps:
+        return ""
+
+
+    class_timestamps = {}
+
+    for timestamp, class_id in zip(timestamps, classes):
+        class_name = class_names.get(class_id, 'Unknown')
+        if class_name not in class_timestamps:
+            class_timestamps[class_name] = []
+        class_timestamps[class_name].append(timestamp)
+
+
+    formatted_ranges = []
+
+    for class_name, timestamps in class_timestamps.items():
+        timestamps = sorted(timestamps)
+        ranges = []
+        start = timestamps[0]
+        for i in range(1, len(timestamps)):
+            if timestamps[i] - timestamps[i - 1] <= 1:
+                continue
+            else:
+                ranges.append(f"{int(start)}-{int(timestamps[i - 1])}")
+                start = timestamps[i]
+
+        ranges.append(f"{int(start)}-{int(timestamps[-1])}")
+
+        formatted_ranges.append(f"{class_name} = {', '.join(ranges)}")
+
+    return ", ".join(formatted_ranges)
+
+import os
+
+def process_video(input_video):
+    cap = cv2.VideoCapture(input_video)
+    if not cap.isOpened():
+        print("Error: Could not open input video.")
+        return None, []
+
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+    # Define the output video path in the current directory
+    output_video_path = os.path.join(os.getcwd(), "processed_video.mp4")
+
+    writer = imageio.get_writer(output_video_path, fps=fps, codec="h264")
+
+    frame_count = 0
+    timestamps = []
+    classes_detected = []
+
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        timestamp = frame_count / fps
+        frame_count += 1
+
+        # Resize the frame to 640x640 before passing to the model
+        resized_frame = cv2.resize(frame, (640, 640))
+
+        gray_frame = cv2.cvtColor(resized_frame, cv2.COLOR_BGR2GRAY)
+        input_frame = cv2.merge([gray_frame, gray_frame, gray_frame])
+
+        results = model.predict(input_frame)
+
+        for result in results:
+            for box in result.boxes:
+                if box.conf[0] >= confidence_threshold:
+                    x1, y1, x2, y2 = map(int, box.xyxy[0])
+                    class_id = int(box.cls[0])
+                    class_name = class_names.get(class_id, f"Class {class_id}")
+                    color = class_colors.get(class_id, (0, 255, 0))
+                    cv2.rectangle(resized_frame, (x1, y1), (x2, y2), color, 2)
+                    text = f'{class_name}, Conf: {box.conf[0]:.2f}'
+                    text_position = (x1, y1 - 10 if y1 > 20 else y1 + 20)
+                    cv2.putText(resized_frame, text, text_position, cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+
+                    timestamps.append(timestamp)
+                    classes_detected.append(class_id)
+
+        # Resize the frame back to original size for the output video
+        output_frame = cv2.resize(resized_frame, (frame_width, frame_height))
+
+        writer.append_data(cv2.cvtColor(output_frame, cv2.COLOR_BGR2RGB))
+
+    cap.release()
+    writer.close()
+
+    formatted_time_ranges = format_time_ranges(timestamps, classes_detected)
+
+    print(f"Processed video saved at: {output_video_path}")
+
+    return output_video_path, formatted_time_ranges
+
+
+
+
+def process_image(input_image):
+    # Convert image from RGB to BGR for OpenCV processing
+    bgr_frame = cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR)
+
+    # Resize the frame to 640x640 before passing to the model
+    resized_frame = cv2.resize(bgr_frame, (640, 640))
+
+    # Convert to grayscale and create a 3-channel grayscale image
+    gray_frame = cv2.cvtColor(resized_frame, cv2.COLOR_BGR2GRAY)
+    input_frame = cv2.merge([gray_frame, gray_frame, gray_frame])
+
+    
+    results = model.predict(input_frame)
+
+    detections_log = []  
+    classes_detected = []  
+
+    for result in results:
+        for box in result.boxes:
+            if box.conf[0] >= confidence_threshold:  
+                x1, y1, x2, y2 = map(int, box.xyxy[0]) 
+                class_id = int(box.cls[0]) 
+                class_name = class_names.get(class_id, f"Class {class_id}")
+                color = class_colors.get(class_id, (0, 255, 0))  # Default green color
+
+               
+                cv2.rectangle(resized_frame, (x1, y1), (x2, y2), color, 2)
+                text = f'{class_name}, Conf: {box.conf[0]:.2f}'
+                text_position = (x1, y1 - 10 if y1 > 20 else y1 + 20)
+                cv2.putText(resized_frame, text, text_position, cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+
+               
+                detections_log.append({
+                    "class": class_name,
+                    "confidence": box.conf[0]
+                })
+                classes_detected.append(class_id)
+
+    # Count occurrences of each class detected
+    class_count = {class_names.get(cls, f"Class {cls}"): classes_detected.count(cls) for cls in set(classes_detected)}
+
+    # Format the detections as 'Class = Count' pairs
+    formatted_log = ", ".join([f"{class_name} = {count}" for class_name, count in class_count.items()])
+
+    # Convert the output frame back to RGB
+    output_image = cv2.cvtColor(resized_frame, cv2.COLOR_BGR2RGB)
+    return output_image, formatted_log
+
+
+with gr.Blocks() as app:
+    gr.Markdown("## Judol Detection using YOLOv11")
+
+    with gr.Tab("Video Detection"):
+        with gr.Row():
+            input_video = gr.Video(label="Upload a video")
+            output_video = gr.Video(label="Processed Video")
+            detections_log = gr.Textbox(label="Detections Log", lines=10)
+
+        input_video.change(
+            fn=lambda input_video: process_video(input_video) if input_video else ("", []),
+            inputs=input_video,
+            outputs=[output_video, detections_log],
+        )
+
+    with gr.Tab("Image Detection"):
+        with gr.Row():
+            input_image = gr.Image(label="Upload an image")
+            output_image = gr.Image(label="Processed Image")
+            image_detections_log = gr.Textbox(label="Detections Log", lines=10)
+
+        input_image.change(
+            fn=process_image,
+            inputs=input_image,
+            outputs=[output_image, image_detections_log],
+        )
+        
+app.launch()