Update app.py

app.py

@@ -1,95 +1,171 @@
-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
-
-# 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)
-
+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)
     return image
 
-def process_video(video_path, target, progress=gr.Progress()):
+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):
     if video_path is None:
-        return None, "Error: No video uploaded"
+        return None, None, None, None, None, None, "Error: No video uploaded"
 
     if not os.path.exists(video_path):
-        return None, f"Error: Video file not found at {video_path}"
+        return None, None, None, None, None, None, f"Error: Video file not found at {video_path}"
 
     cap = cv2.VideoCapture(video_path)
     if not cap.isOpened():
-        return None, f"Error: Unable to open video file at {video_path}"
+        return None, None, None, None, 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))
+    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
     original_duration = frame_count / original_fps
-    output_fps = 5
+    
+    frame_interval = max(1, round(original_fps / 10))  # Process 10 frames per second
+    
+    body_movements = []
+    time_points = []
 
-    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
+    prev_poses = None
     frames = []
+    frame_indices = []
 
-    for frame in progress.tqdm(range(frame_count)):
+    for frame in progress.tqdm(range(0, frame_count, frame_interval)):
+        cap.set(cv2.CAP_PROP_POS_FRAMES, frame)
         ret, img = cap.read()
         if not ret:
             break
+        frames.append(img)
+        frame_indices.append(frame)
 
-        if frame % (original_fps // output_fps) != 0:
-            continue
+        if len(frames) == batch_size:
+            process_batch(frames, frame_indices, prev_poses, body_movements, time_points, original_fps)
+            frames = []
 
-        pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
-        frames.append(pil_img)
+    # Process any remaining frames
+    if frames:
+        process_batch(frames, frame_indices, prev_poses, body_movements, time_points, original_fps)
 
-        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 = []
+    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 not os.path.exists(video_path):
+        return None
 
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        return None
+
+    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()
 
-    return output_path, None
+    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
 
 def load_sample_frame(video_path):
     cap = cv2.VideoCapture(video_path)
@@ -104,35 +180,57 @@ def load_sample_frame(video_path):
 
 def gradio_app():
     with gr.Blocks() as app:
-        gr.Markdown("# Video Object Detection with Owlv2")
-
+        gr.Markdown("# Multi-Person Body Movement Analysis")
+        
         video_input = gr.Video(label="Upload Video")
-        target_input = gr.Textbox(label="Target Object")
-        output_video = gr.Video(label="Output 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")
+        
         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")
-
+        
         video_path = gr.State(None)
-        def process_and_update(video, target):
-            output_video_path, error = process_video(video, target)
-            if error:
-                error_output.visible = True
-            else:
+        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
-            return output_video_path, error
+            else:
+                error_output.visible = True
+            return fig, video, error, video_path_value, fig, ax, time_points, body_movements, video_duration
 
         video_input.upload(process_and_update, 
-                           inputs=[video_input, target_input], 
-                           outputs=[output_video, error_output])
+                           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])
 
         def use_sample_video():
-            sample_video_path = "Drone Video of African Wildlife Wild Botswan.mp4"
-            return process_and_update(sample_video_path, "animal")
+            sample_video_path = "IL_Dancing_Sample.mp4"
+            return process_and_update(sample_video_path)
 
         use_sample_button.click(use_sample_video, 
                                 inputs=None, 
-                                outputs=[output_video, error_output])
+                                outputs=[graph_output, video_output, error_output, video_path, 
+                                         fig_state, ax_state, time_points_state, body_movements_state, video_duration_state])
 
     return app