app.py

import spaces
import gradio as gr
import cv2
from PIL import Image, ImageDraw, ImageFont
import torch
from transformers import Owlv2Processor, Owlv2ForObjectDetection
import numpy as np
import os
import matplotlib.pyplot as plt
import tempfile
import shutil

device = "cuda"

processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16")
model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16")

model = model.to(device)

def process_video(video_path, target, progress=gr.Progress()):
    if video_path is None:
        return None, None, "Error: No video uploaded"

    if not os.path.exists(video_path):
        return None, None, f"Error: Video file not found at {video_path}"

    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return None, None, f"Error: Unable to open video file at {video_path}"

    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    original_fps = int(cap.get(cv2.CAP_PROP_FPS))
    output_fps = 1
    frame_duration = 1 / output_fps
    video_duration = frame_count / original_fps

    frame_scores = []
    temp_dir = tempfile.mkdtemp()
    frame_paths = []

    batch_size = 1
    batch_frames = []
    batch_indices = []

    for i, time in enumerate(progress.tqdm(np.arange(0, video_duration, frame_duration))):
        frame_number = int(time * original_fps)
        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
        ret, img = cap.read()
        if not ret:
            break

        # Convert to RGB without resizing
        pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        
        batch_frames.append(pil_img)
        batch_indices.append(i)

        if len(batch_frames) == batch_size or i == int(video_duration / frame_duration) - 1:
            # Process batch
            inputs = processor(text=[target] * len(batch_frames), images=batch_frames, return_tensors="pt", padding=True).to(device)
            
            with torch.no_grad():
                outputs = model(**inputs)

            target_sizes = torch.Tensor([pil_img.size[::-1] for _ in batch_frames]).to(device)
            results = processor.post_process_object_detection(outputs, target_sizes=target_sizes)

            for idx, (pil_img, result) in enumerate(zip(batch_frames, results)):
                draw = ImageDraw.Draw(pil_img)
                max_score = 0

                boxes, scores, labels = result["boxes"], result["scores"], result["labels"]

                # Inside the loop where bounding boxes are drawn
                for box, score, label in zip(boxes, scores, labels):
                    if score.item() >= 0.5:
                        box = [round(i, 2) for i in box.tolist()]
                        object_label = target
                        confidence = round(score.item(), 3)
                        annotation = f"{object_label}: {confidence}"
                
                        # Increase line width for the bounding box
                        draw.rectangle(box, outline="red", width=3)
                
                        # Calculate font size based on image dimensions
                        img_width, img_height = pil_img.size
                        font_size = int(min(img_width, img_height) * 0.03)  # 3% of the smaller dimension
                        try:
                            font = ImageFont.truetype("arial.ttf", font_size)
                        except IOError:
                            font = ImageFont.load_default()
                
                        # Calculate text size
                        text_bbox = draw.textbbox((0, 0), annotation, font=font)
                        text_width = text_bbox[2] - text_bbox[0]
                        text_height = text_bbox[3] - text_bbox[1]
                
                        # Position text inside the top of the bounding box
                        text_position = (box[0], box[1])
                
                        # Draw semi-transparent background for text
                        draw.rectangle([text_position[0], text_position[1], 
                                        text_position[0] + text_width, text_position[1] + text_height], 
                                       fill=(0, 0, 0, 128))
                
                        # Draw text in red
                        draw.text(text_position, annotation, fill="red", font=font)
                
                        max_score = max(max_score, confidence)

                frame_path = os.path.join(temp_dir, f"frame_{batch_indices[idx]:04d}.png")
                pil_img.save(frame_path)
                frame_paths.append(frame_path)
                frame_scores.append(max_score)

            # Clear batch
            batch_frames = []
            batch_indices = []

        # Clear GPU cache every 10 frames
        if i % 10 == 0:
            torch.cuda.empty_cache()

    cap.release()
    return frame_paths, frame_scores, None

def create_heatmap(frame_scores, current_frame):
    plt.figure(figsize=(16, 4))
    plt.imshow([frame_scores], cmap='hot_r', aspect='auto')
    plt.title('Object Detection Heatmap', fontsize=14)
    plt.xlabel('Frame', fontsize=12)
    plt.yticks([])

    num_frames = len(frame_scores)
    step = max(1, num_frames // 20)
    frame_numbers = range(0, num_frames, step)
    plt.xticks(frame_numbers, [str(i) for i in frame_numbers], rotation=90, ha='right')

    plt.axvline(x=current_frame, color='blue', linestyle='--', linewidth=2)

    plt.tight_layout()
    
    with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as tmp_file:
        plt.savefig(tmp_file.name, format='png', dpi=400, bbox_inches='tight')
    plt.close()
    
    return tmp_file.name

def load_sample_frame(video_path, target_frame=87, original_fps=30, processing_fps=1):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return None
    
    # Calculate the corresponding frame number in the original video
    original_frame_number = int(target_frame * (original_fps / processing_fps))
    
    # Set the frame position
    cap.set(cv2.CAP_PROP_POS_FRAMES, original_frame_number)
    
    ret, frame = cap.read()
    cap.release()
    if not ret:
        return None
    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    return frame_rgb

def update_frame_and_heatmap(frame_index, frame_paths, scores):
    if frame_paths and 0 <= frame_index < len(frame_paths):
        frame = Image.open(frame_paths[frame_index])
        heatmap_path = create_heatmap(scores, frame_index)
        return np.array(frame), heatmap_path
    return None, None

def gradio_app():
    with gr.Blocks() as app:
        gr.Markdown("# Video Object Detection with Owlv2")

        video_input = gr.Video(label="Upload Video")
        target_input = gr.Textbox(label="Target Object", value="Elephant")
        frame_slider = gr.Slider(minimum=0, maximum=100, step=1, label="Frame", value=0)
        heatmap_output = gr.Image(label="Detection Heatmap")
        output_image = gr.Image(label="Processed Frame")
        error_output = gr.Textbox(label="Error Messages", visible=False)
        sample_video_frame = gr.Image(
            value=load_sample_frame("Drone Video of African Wildlife Wild Botswan.mp4", target_frame=87),
            label="Drone Video of African Wildlife Wild Botswan by wildimagesonline.com - Sample Video Frame (Frame 87 at 1 FPS)"
        )
        use_sample_button = gr.Button("Use Sample Video")
        progress_bar = gr.Progress()

        frame_paths = gr.State([])
        frame_scores = gr.State([])

        def process_and_update(video, target):
            paths, scores, error = process_video(video, target, progress_bar)
            if paths is not None:
                heatmap_path = create_heatmap(scores, 0)
                first_frame = Image.open(paths[0])
                return paths, scores, np.array(first_frame), heatmap_path, error, gr.Slider(maximum=len(paths) - 1, value=0)
            return None, None, None, None, error, gr.Slider(maximum=100, value=0)

        video_input.upload(process_and_update, 
                           inputs=[video_input, target_input], 
                           outputs=[frame_paths, frame_scores, output_image, heatmap_output, error_output, frame_slider])

        frame_slider.change(update_frame_and_heatmap, 
                            inputs=[frame_slider, frame_paths, frame_scores], 
                            outputs=[output_image, heatmap_output])

        def use_sample_video():
            sample_video_path = "Drone Video of African Wildlife Wild Botswan.mp4"
            return process_and_update(sample_video_path, "Elephant")

        use_sample_button.click(use_sample_video, 
                                inputs=None, 
                                outputs=[frame_paths, frame_scores, output_image, heatmap_output, error_output, frame_slider])

        # Layout
        with gr.Row():
            with gr.Column(scale=2):
                output_image
            with gr.Column(scale=1):
                sample_video_frame
                use_sample_button

    return app

if __name__ == "__main__":
    app = gradio_app()
    app.launch()

    # Cleanup temporary files
    def cleanup():
        for path in frame_paths.value:
            if os.path.exists(path):
                os.remove(path)
        if os.path.exists(temp_dir):
            shutil.rmtree(temp_dir)

    # Make sure to call cleanup when the app is closed
    # This might require additional setup depending on how you're running the app