import torch
from ultralytics import YOLO
import gradio as gr
from PIL import Image
import numpy as np

class ObjectDetectorV8:
    def __init__(self):
        # Load YOLOv8 model
        self.model = YOLO('yolov8n.pt')  
        
    def count_objects(self, image):
        """
        Detect and count objects in the input image.
        """
        # Pastikan gambar tidak None
        if image is None:
            return None, {}
        
        # Convert PIL Image to numpy array jika perlu
        image_np = np.array(image)
        
        # Perform object detection
        results = self.model(image_np)
        
        # Extract results
        detections = results[0].boxes.data.cpu().numpy()
        labels = detections[:, 5].astype(int)
        
        # Convert labels to human-readable names dan count occurrences
        label_names = [self.model.names[label] for label in labels]
        count_dict = {label: label_names.count(label) for label in set(label_names)}
        
        # Annotate image with bounding boxes
        annotated_image = results[0].plot()  
        
        return Image.fromarray(annotated_image), count_dict
    
    def live_object_detection(self, video):
        """
        Perform live object detection on video stream
        """
        # Pastikan video tidak None
        if video is None:
            return None, {}
        
        # Perform object detection
        results = self.model(video)
        
        # Extract results
        detections = results[0].boxes.data.cpu().numpy()
        labels = detections[:, 5].astype(int)
        
        # Convert labels to human-readable names dan count occurrences
        label_names = [self.model.names[label] for label in labels]
        count_dict = {label: label_names.count(label) for label in set(label_names)}
        
        # Annotate image with bounding boxes
        annotated_image = results[0].plot()
        
        return annotated_image, count_dict

# Inisialisasi detector sekali
detector = ObjectDetectorV8()

# Buat antarmuka Gradio
def create_interface():
    image_interface = gr.Interface(
        fn=detector.count_objects,
        inputs=gr.Image(type="pil", label="Upload Gambar"),
        outputs=[
            gr.Image(label="Objek Terdeteksi"),
            gr.JSON(label="Jumlah Objek")
        ],
        title="Pendeteksi Objek dengan YOLOv8",
        description="Unggah gambar untuk mendeteksi dan menghitung objek"
    )
    
    video_interface = gr.Interface(
        fn=detector.live_object_detection,
        inputs=gr.Image(sources=["webcam"], type="numpy", label="Input Webcam"),
        outputs=[
            gr.Image(label="Deteksi Langsung"),
            gr.JSON(label="Jumlah Objek")
        ],
        title="Deteksi Objek Langsung",
        description="Gunakan webcam untuk deteksi objek real-time"
    )
    
    demo = gr.TabbedInterface(
        [image_interface, video_interface], 
        ["Deteksi Gambar", "Deteksi Langsung"]
    )
    
    return demo

# Buat dan luncurkan interface
iface = create_interface()

if __name__ == "__main__":
    iface.launch()