med333gg/fruit-detection-api

1import gradio as gr2import cv23import numpy as np4from ultralytics import YOLO5import json6import os7import tempfile8 9# Load YOLOv8 model (will run on GPU if available)10model = YOLO('yolov8n.pt')11 12def process_video(video_path):13    """Process video and return detection results"""14    if video_path is None:15        return None, "No video uploaded"16    17    try:18        # Create temporary directory for results19        with tempfile.TemporaryDirectory() as temp_dir:20            output_path = os.path.join(temp_dir, 'output.mp4')21            22            # Open video23            cap = cv2.VideoCapture(video_path)24            if not cap.isOpened():25                return None, "Error: Could not open video file"26            27            fps = cap.get(cv2.CAP_PROP_FPS)28            w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))29            h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))30            31            # Video writer32            fourcc = cv2.VideoWriter_fourcc(*'mp4v')33            out = cv2.VideoWriter(output_path, fourcc, fps, (w, h))34            35            total_strawberries = 036            frame_idx = 037            38            # Process frames (limit to first 100 frames for demo)39            max_frames = 10040            41            while frame_idx < max_frames:42                ret, frame = cap.read()43                if not ret:44                    break45                46                # Run detection on GPU47                results = model(frame)48                49                strawberries = 050                for box in results[0].boxes:51                    cls = int(box.cls[0])52                    # COCO class 53 is 'apple', treat as 'strawberry'53                    if cls == 53:54                        strawberries += 155                        xyxy = box.xyxy[0].cpu().numpy().astype(int)56                        cv2.rectangle(frame, (xyxy[0], xyxy[1]), (xyxy[2], xyxy[3]), (0, 255, 0), 2)57                        cv2.putText(frame, 'Strawberry', (xyxy[0], xyxy[1] - 10), 58                                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)59                60                total_strawberries += strawberries61                out.write(frame)62                frame_idx += 163            64            cap.release()65            out.release()66            67            return output_path, f"Detected {total_strawberries} strawberries in {frame_idx} frames"68    69    except Exception as e:70        return None, f"Error processing video: {str(e)}"71 72# Create Gradio interface73with gr.Blocks(title="Fruit Detection API") as demo:74    gr.Markdown("# 🍓 Fruit Detection API")75    gr.Markdown("Upload a video to detect strawberries using YOLOv8 on GPU")76    gr.Markdown("**Note:** Processing is limited to first 100 frames for demo purposes")77    78    with gr.Row():79        video_input = gr.Video(label="Upload Video")80        video_output = gr.Video(label="Processed Video")81    82    text_output = gr.Textbox(label="Results")83    84    process_btn = gr.Button("Process Video", variant="primary")85    process_btn.click(86        fn=process_video,87        inputs=[video_input],88        outputs=[video_output, text_output]89    )90 91# Launch the app92if __name__ == "__main__":93    demo.launch()