Update app.py

app.py

@@ -1,171 +1,95 @@
+import gradio as gr
 import cv2
+from PIL import Image, ImageDraw, ImageFont
+import torch
+from transformers import Owlv2Processor, Owlv2ForObjectDetection
 import numpy as np
-import matplotlib.pyplot as plt
-import gradio as gr
-from moviepy.editor import *
 import os
-import torch
-import openpifpaf
-
-# Ensure NumPy is available
-try:
-    import numpy as np
-except ImportError:
-    os.system('pip install numpy')
-    import numpy as np
-
-# OpenPifPaf configuration
-predictor = openpifpaf.Predictor(checkpoint='shufflenetv2k16')
-
-def preprocess(image):
-    input_size = (192, 256)
-    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-    image = cv2.resize(image, input_size)
+
+# Check if CUDA is available, otherwise use CPU
+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 detect_objects_in_frame(image, target):
+    draw = ImageDraw.Draw(image)
+    texts = [[target]]
+    inputs = processor(text=texts, images=image, return_tensors="pt", padding=True).to(device)
+    outputs = model(**inputs)
+
+    target_sizes = torch.Tensor([image.size[::-1]])
+    results = processor.post_process_object_detection(outputs=outputs, threshold=0.1, target_sizes=target_sizes)
+
+    color_map = {target: "red"}
+
+    try:
+        font = ImageFont.truetype("arial.ttf", 15)
+    except IOError:
+        font = ImageFont.load_default()
+
+    i = 0
+    text = texts[i]
+    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.25:
+            box = [round(i, 2) for i in box.tolist()]
+            object_label = text[label]
+            confidence = round(score.item(), 3)
+            annotation = f"{object_label}: {confidence}"
+
+            draw.rectangle(box, outline=color_map.get(object_label, "red"), width=2)
+            text_position = (box[0], box[1] - 10)
+            draw.text(text_position, annotation, fill="white", font=font)
+
     return image
 
-def total_body_movement(current_poses, prev_poses):
-    if not current_poses or not prev_poses:
-        return 0
-    total_movement = 0
-    for current_pose in current_poses:
-        for prev_pose in prev_poses:
-            movement = np.sum(np.sqrt(np.sum((current_pose - prev_pose)**2, axis=1)))
-            total_movement += movement
-    return total_movement / (len(current_poses) * len(prev_poses))
-
-def process_video(video_path, progress=gr.Progress(), batch_size=64):
+def process_video(video_path, target, progress=gr.Progress()):
     if video_path is None:
-        return None, None, None, None, None, None, "Error: No video uploaded"
+        return None, "Error: No video uploaded"
 
     if not os.path.exists(video_path):
-        return None, None, None, None, None, None, f"Error: Video file not found at {video_path}"
+        return None, f"Error: Video file not found at {video_path}"
 
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
-        return None, None, None, None, None, None, f"Error: Unable to open video file at {video_path}"
+        return None, f"Error: Unable to open video file at {video_path}"
 
-    original_fps = int(cap.get(cv2.CAP_PROP_FPS))
     frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    original_fps = int(cap.get(cv2.CAP_PROP_FPS))
     original_duration = frame_count / original_fps
-    
-    frame_interval = max(1, round(original_fps / 10))  # Process 10 frames per second
-    
-    body_movements = []
-    time_points = []
+    output_fps = 5
 
-    prev_poses = None
+    output_path = "output_video.mp4"
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(output_path, fourcc, output_fps, (int(cap.get(3)), int(cap.get(4))))
+
+    batch_size = 64
     frames = []
-    frame_indices = []
 
-    for frame in progress.tqdm(range(0, frame_count, frame_interval)):
-        cap.set(cv2.CAP_PROP_POS_FRAMES, frame)
+    for frame in progress.tqdm(range(frame_count)):
         ret, img = cap.read()
         if not ret:
             break
-        frames.append(img)
-        frame_indices.append(frame)
-
-        if len(frames) == batch_size:
-            process_batch(frames, frame_indices, prev_poses, body_movements, time_points, original_fps)
-            frames = []
-
-    # Process any remaining frames
-    if frames:
-        process_batch(frames, frame_indices, prev_poses, body_movements, time_points, original_fps)
-
-    cap.release()
 
-    fig, ax = plt.subplots(figsize=(10, 6), dpi=500)
-    ax.plot(time_points, body_movements, "-", linewidth=0.5)
-    ax.set_xlim(0, original_duration)
-    ax.set_xlabel("Time")
-    ax.set_ylabel("Body Movement")
-    ax.set_title("Body Movement Analysis")
-
-    num_labels = 50
-    label_positions = np.linspace(0, original_duration, num_labels)
-    label_texts = [f"{int(t//60):02d}:{int(t%60):02d}" for t in label_positions]
-    ax.set_xticks(label_positions)
-    ax.set_xticklabels(label_texts, rotation=90, ha='right')
-    plt.tight_layout()
-
-    return fig, ax, time_points, body_movements, video_path, original_duration, None
-
-def process_batch(frames, frame_indices, prev_poses, body_movements, time_points, original_fps):
-    batch_preds = predictor.numpy_images(frames)
-
-    for i, (predictions, frame_index) in enumerate(zip(batch_preds, frame_indices)):
-        pose_coords = [pred.data for pred in predictions]
-
-        if prev_poses is not None:
-            movement = total_body_movement(pose_coords, prev_poses)
-            body_movements.append(movement)
-        else:
-            body_movements.append(0)
-        
-        prev_poses = pose_coords
-        time_points.append(frame_index / original_fps)
-
-def update_video(video_path, time):
-    if video_path is None:
-        return None
+        if frame % (original_fps // output_fps) != 0:
+            continue
 
-    if not os.path.exists(video_path):
-        return None
+        pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+        frames.append(pil_img)
 
-    cap = cv2.VideoCapture(video_path)
-    if not cap.isOpened():
-        return None
+        if len(frames) == batch_size or frame == frame_count - 1:
+            annotated_frames = [detect_objects_in_frame(frame, target) for frame in frames]
+            for annotated_img in annotated_frames:
+                annotated_frame = cv2.cvtColor(np.array(annotated_img), cv2.COLOR_RGB2BGR)
+                out.write(annotated_frame)
+            frames = []
 
-    original_fps = int(cap.get(cv2.CAP_PROP_FPS))
-    frame_number = int(time * original_fps)
-    cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
-    ret, img = cap.read()
     cap.release()
+    out.release()
 
-    if not ret:
-        return None
-
-    predictions, _, _ = predictor.numpy_image(img)
-    pose_coords = [pred.data for pred in predictions]
-
-    for coords in pose_coords:
-        for i in range(len(coords)):
-            x, y = coords[i]
-            if x > 0 and y > 0:
-                cv2.circle(img, (int(x), int(y)), 3, (0, 255, 0), -1)
-    
-    for pred in predictions:
-        skeleton = pred.data[:, :2]
-        for i, j in pred.skeleton:
-            if skeleton[i, 0] > 0 and skeleton[i, 1] > 0 and skeleton[j, 0] > 0 and skeleton[j, 1] > 0:
-                cv2.line(img, (int(skeleton[i, 0]), int(skeleton[i, 1])), (int(skeleton[j, 0]), int(skeleton[j, 1])), (255, 0, 0), 2)
-
-    return img
-
-def update_graph(fig, ax, time_points, body_movements, current_time, video_duration):
-    ax.clear()
-    ax.plot(time_points, body_movements, "-", linewidth=0.5)
-    ax.axvline(x=current_time, color='r', linestyle='--')
-    
-    minutes, seconds = divmod(int(current_time), 60)
-    timecode = f"{minutes:02d}:{seconds:02d}"
-    ax.text(current_time, ax.get_ylim()[1], timecode, 
-            verticalalignment='top', horizontalalignment='right',
-            color='r', fontweight='bold', bbox=dict(facecolor='white', edgecolor='none', alpha=0.7))
-    
-    ax.set_xlabel("Time")
-    ax.set_ylabel("Body Movement")
-    ax.set_title("Body Movement Analysis")
-    
-    num_labels = 80
-    label_positions = np.linspace(0, video_duration, num_labels)
-    label_texts = [f"{int(t//60):02d}:{int(t%60):02d}" for t in label_positions]
-    ax.set_xticks(label_positions)
-    ax.set_xticklabels(label_texts, rotation=90, ha='right')
-    ax.set_xlim(0, video_duration)
-    plt.tight_layout()
-    return fig
+    return output_path, None
 
 def load_sample_frame(video_path):
     cap = cv2.VideoCapture(video_path)
@@ -180,57 +104,36 @@ def load_sample_frame(video_path):
 
 def gradio_app():
     with gr.Blocks() as app:
-        gr.Markdown("# Multi-Person Body Movement Analysis")
-        
+        gr.Markdown("# Video Object Detection with Owlv2")
+
         video_input = gr.Video(label="Upload Video")
-        graph_output = gr.Plot()
-        time_slider = gr.Slider(label="Time (seconds)", minimum=0, maximum=100, step=0.1)
-        video_output = gr.Image(label="Body Posture")
-        
-        with gr.Row():
-            sample_video_frame = gr.Image(value=load_sample_frame("IL_Dancing_Sample.mp4"), label="Sample Video Frame")
-            use_sample_button = gr.Button("Use Sample Video")
-        
+        target_input = gr.Textbox(label="Target Object")
+        output_video = gr.Video(label="Output Video")
         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()
+
         video_path = gr.State(None)
-        fig_state = gr.State(None)
-        ax_state = gr.State(None)
-        time_points_state = gr.State(None)
-        body_movements_state = gr.State(None)
-        video_duration_state = gr.State(None)
-
-        def process_and_update(video):
-            fig, ax, time_points, body_movements, video_path_value, video_duration, error = process_video(video)
-            if fig is not None:
-                time_slider.maximum = video_duration
-                error_output.visible = False
-            else:
+        def process_and_update(video, target):
+            output_video_path, error = process_video(video, target, progress_bar)
+            if error:
                 error_output.visible = True
-            return fig, video, error, video_path_value, fig, ax, time_points, body_movements, video_duration
+            else:
+                error_output.visible = False
+            return output_video_path, error
 
         video_input.upload(process_and_update, 
-                           inputs=video_input, 
-                           outputs=[graph_output, video_output, error_output, video_path, 
-                                    fig_state, ax_state, time_points_state, body_movements_state, video_duration_state])
-
-        def update_video_and_graph(video_path_value, current_time, fig, ax, time_points, body_movements, video_duration):
-            updated_frame = update_video(video_path_value, current_time)
-            updated_fig = update_graph(fig, ax, time_points, body_movements, current_time, video_duration)
-            return updated_frame, updated_fig
-
-        time_slider.change(update_video_and_graph, 
-                           inputs=[video_path, time_slider, fig_state, ax_state, time_points_state, body_movements_state, video_duration_state], 
-                           outputs=[video_output, graph_output])
+                           inputs=[video_input, target_input], 
+                           outputs=[output_video, error_output])
 
         def use_sample_video():
-            sample_video_path = "IL_Dancing_Sample.mp4"
-            return process_and_update(sample_video_path)
+            sample_video_path = "Drone Video of African Wildlife Wild Botswan.mp4"
+            return process_and_update(sample_video_path, "animal")
 
         use_sample_button.click(use_sample_video, 
                                 inputs=None, 
-                                outputs=[graph_output, video_output, error_output, video_path, 
-                                         fig_state, ax_state, time_points_state, body_movements_state, video_duration_state])
+                                outputs=[output_video, error_output])
 
     return app