app.py

import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image
import tempfile

model = YOLO("best.pt")

def predict_image(img):
    results = model(img)
    annotated = results[0].plot()
    
    # Hitung jumlah objek
    num_objects = len(results[0].boxes) if results[0].boxes is not None else 0
    info = f"Jumlah sawit terdeteksi: {num_objects}"
    
    return Image.fromarray(annotated), info

def predict_video(video_path):
    if video_path is None:
        return None, "Tidak ada video"
    
    cap = cv2.VideoCapture(video_path)
    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))
    
    # Buat file output
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
    output_path = temp_file.name
    temp_file.close()
    
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
    
    total_detections = 0
    frame_count = 0
    detections_per_frame = []
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
            
        results = model(frame)
        annotated = results[0].plot()
        
        # Hitung deteksi per frame
        frame_detections = len(results[0].boxes) if results[0].boxes is not None else 0
        detections_per_frame.append(frame_detections)
        total_detections += frame_detections
        
        # Tambahkan text info di video
        cv2.putText(annotated, f"Frame {frame_count+1}: {frame_detections} objek", 
                   (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
        
        out.write(annotated)
        frame_count += 1
    
    cap.release()
    out.release()
    
    # Buat info detail
    avg_detections = total_detections / frame_count if frame_count > 0 else 0
    max_detections = max(detections_per_frame) if detections_per_frame else 0
    
    info = f"""Total frame: {frame_count}
Total deteksi: {total_detections}
Rata-rata per frame: {avg_detections:.1f}
Maksimum per frame: {max_detections}
Detail per frame: {detections_per_frame[:10]}{'...' if len(detections_per_frame) > 10 else ''}"""
    
    return output_path, info

# Interface
image_demo = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Image(type="pil"), gr.Textbox(label="Info")],
    title="Deteksi Sawit - Gambar"
)

video_demo = gr.Interface(
    fn=predict_video,
    inputs=gr.Video(),
    outputs=[gr.Video(), gr.Textbox(label="Info")],
    title="Deteksi Sawit - Video"
)

gr.TabbedInterface([image_demo, video_demo], ["Gambar", "Video"]).launch(share=True)