|
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 |
|
from io import BytesIO |
|
import base64 |
|
|
|
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu' |
|
|
|
processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16") |
|
model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16").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 = 3 |
|
frame_duration = 1 / output_fps |
|
video_duration = frame_count / original_fps |
|
|
|
processed_frames = [] |
|
frame_scores = [] |
|
batch_size = 4 |
|
batch_frames = [] |
|
batch_times = [] |
|
|
|
for time in 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 |
|
|
|
pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) |
|
batch_frames.append(pil_img) |
|
batch_times.append(time) |
|
|
|
if len(batch_frames) == batch_size or time + frame_duration >= video_duration: |
|
|
|
texts = [[target]] * len(batch_frames) |
|
inputs = processor(text=texts, images=batch_frames, return_tensors="pt", padding=True).to(device) |
|
outputs = model(**inputs) |
|
|
|
for i, (image, batch_time) in enumerate(zip(batch_frames, batch_times)): |
|
target_sizes = torch.Tensor([image.size[::-1]]) |
|
results = processor.post_process_object_detection(outputs, target_sizes=target_sizes) |
|
|
|
draw = ImageDraw.Draw(image) |
|
max_score = 0 |
|
|
|
try: |
|
font = ImageFont.truetype("arial.ttf", 30) |
|
except IOError: |
|
font = ImageFont.load_default() |
|
|
|
boxes, scores, labels = results[i]["boxes"], results[i]["scores"], results[i]["labels"] |
|
|
|
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}" |
|
|
|
draw.rectangle(box, outline="red", width=4) |
|
text_position = (box[0], box[1] - 30) |
|
draw.text(text_position, annotation, fill="white", font=font) |
|
|
|
max_score = max(max_score, confidence) |
|
|
|
processed_frames.append(np.array(image)) |
|
frame_scores.append(max_score) |
|
|
|
batch_frames = [] |
|
batch_times = [] |
|
|
|
cap.release() |
|
return processed_frames, frame_scores, None |
|
|
|
def create_heatmap(frame_scores): |
|
plt.figure(figsize=(10, 2)) |
|
plt.imshow([frame_scores], cmap='hot', aspect='auto') |
|
plt.colorbar(label='Confidence') |
|
plt.title('Object Detection Heatmap') |
|
plt.xlabel('Frame') |
|
plt.yticks([]) |
|
plt.tight_layout() |
|
|
|
buf = BytesIO() |
|
plt.savefig(buf, format='png') |
|
buf.seek(0) |
|
plt.close() |
|
|
|
return base64.b64encode(buf.getvalue()).decode('utf-8') |
|
|
|
def load_sample_frame(video_path): |
|
cap = cv2.VideoCapture(video_path) |
|
if not cap.isOpened(): |
|
return None |
|
ret, frame = cap.read() |
|
cap.release() |
|
if not ret: |
|
return None |
|
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) |
|
return frame_rgb |
|
|
|
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) |
|
output_image = gr.Image(label="Processed Frame") |
|
heatmap_output = gr.Image(label="Detection Heatmap") |
|
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"), label="Sample Video Frame") |
|
use_sample_button = gr.Button("Use Sample Video") |
|
progress_bar = gr.Progress() |
|
|
|
processed_frames = gr.State([]) |
|
frame_scores = gr.State([]) |
|
|
|
def process_and_update(video, target): |
|
frames, scores, error = process_video(video, target, progress_bar) |
|
if frames is not None: |
|
heatmap = create_heatmap(scores) |
|
return frames, scores, frames[0], heatmap, error, gr.Slider(maximum=len(frames) - 1, value=0) |
|
return None, None, None, None, error, gr.Slider(maximum=100, value=0) |
|
|
|
def update_frame(frame_index, frames): |
|
if frames and 0 <= frame_index < len(frames): |
|
return frames[frame_index] |
|
return None |
|
|
|
video_input.upload(process_and_update, |
|
inputs=[video_input, target_input], |
|
outputs=[processed_frames, frame_scores, output_image, heatmap_output, error_output, frame_slider]) |
|
|
|
frame_slider.change(update_frame, |
|
inputs=[frame_slider, processed_frames], |
|
outputs=[output_image]) |
|
|
|
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=[processed_frames, frame_scores, output_image, heatmap_output, error_output, frame_slider]) |
|
|
|
return app |
|
|
|
if __name__ == "__main__": |
|
app = gradio_app() |
|
app.launch(share=True) |
