"""Helpers for visual prompts.""" import os import random import numpy as np import cv2 import decord from PIL import Image, ImageDraw, ImageFont def add_absolute_time_bar(video_path, save_path, bar_height=20, bar_color="red", max_duration=10.): from moviepy.editor import VideoFileClip from moviepy.video.fx.all import crop # Load the video clip video_clip = VideoFileClip(video_path) # Get video properties video_duration = video_clip.duration # Total duration of the video in seconds video_width, video_height = video_clip.size # Convert bar color (e.g., "red") to RGB color_map = {"red": (255, 0, 0), "green": (0, 255, 0), "blue": (0, 0, 255)} # Extendable bar_color_rgb = color_map.get(bar_color.lower(), (255, 0, 0)) # Default is red def add_bar_to_frame(get_frame, t): """ Function to add a time bar to each frame based on time t (in seconds). """ frame = get_frame(t) current_time = min(t, max_duration) # Cap the time at max_duration # Calculate the proportional bar length bar_length = int((current_time / max_duration) * video_width) # Create a new bar image with the same width as the frame bar = np.zeros((bar_height, video_width, 3), dtype=np.uint8) bar[:, :bar_length] = bar_color_rgb # Stack the bar on top of the current frame frame_with_bar = np.vstack([bar, frame]) return frame_with_bar # Modify the video frames by adding the time bar video_with_bar = video_clip.fl(add_bar_to_frame) # Crop the extra height from the video (if needed) video_with_bar = crop(video_with_bar, height=video_height) # Keeps original video height # Save the modified video to the output path video_with_bar.write_videofile(save_path, codec="libx264", verbose=False, logger=None) def add_clock(video_path, save_path, clock_color='red', clock_location=None, clock_radius=25): from moviepy.editor import VideoFileClip from moviepy.video.fx.all import crop # Load the video clip video_clip = VideoFileClip(video_path) # Get video properties video_width, video_height = video_clip.size video_duration = video_clip.duration # Total duration of the video in seconds # Default clock location if not provided: 10% height and width from the top-right corner if clock_location is None: clock_location = [int(video_width * 0.9), int(video_height * 0.1)] # Make sure the clock does not go out of the image clock_location[0] = min(clock_location[0], video_width - clock_radius) clock_location[1] = max(clock_location[1], clock_radius) clock_location = tuple(clock_location) # Convert clock_color from string to BGR for OpenCV color_map = {"red": (255, 0, 0), "green": (0, 255, 0), "blue": (0, 0, 255)} # Extendable bar_color_rgb = color_map.get(clock_color.lower(), (255, 0, 0)) # Default is red def add_clock_to_frame(get_frame, t): """ Function to add a clock bar (circle) that fills clockwise as time progresses. """ frame = get_frame(t) # Create a copy of the frame to modify frame = np.array(frame) # Calculate the angle based on the current time t relative to the total video duration angle = int((t / video_duration) * 360) # From 0 to 360 degrees # Draw a partial filled circle (arc) to represent the clock progress center = clock_location # The center of the clock thickness = -1 # Solid fill # Draw the circular progress using cv2.ellipse cv2.ellipse( frame, center, (clock_radius, clock_radius), # Clock radius in both x and y direction -90, # Start angle 0, # Starting point (0 degrees) angle, # End angle based on progress bar_color_rgb, # Clock color in BGR thickness # Fill the shape ) return frame # Apply the clock bar to each frame in the video video_with_clock = video_clip.fl(add_clock_to_frame) # Save the modified video with the clock to the output path video_with_clock.write_videofile(save_path, codec="libx264", verbose=False, logger=None) def add_expanding_shape( video_path, save_path, shape='circle', max_duration=None, shape_color='blue', shape_location=None, min_size=10, max_size=50, thickness=3, ): from moviepy.editor import VideoFileClip from moviepy.video.fx.all import crop # Load the video clip video_clip = VideoFileClip(video_path) # Get video properties video_width, video_height = video_clip.size video_duration = video_clip.duration if max_duration is None else max_duration # Default shape location if not provided: 10% width and height from the top-right corner if shape_location is None: shape_location = [int(video_width * 0.9), int(video_height * 0.1)] # # Make sure the shape does not go out of the image # shape_location[0] = min(shape_location[0], video_width - max_size) # shape_location[1] = max(shape_location[1], max_size) # Move the location if the shape goes out of the image if shape_location[0] + max_size > video_width: shape_location[0] = video_width - max_size if shape_location[1] - max_size < 0: shape_location[1] = max_size shape_location = tuple(shape_location) else: if isinstance(shape_location, str): shape_location = shape_location.lower() if shape_location == 'center': shape_location = (video_width // 2, video_height // 2) elif shape_location == 'top-left': shape_location = (max_size, max_size) elif shape_location == 'top-right': shape_location = (video_width - max_size, max_size) elif shape_location == 'bottom-left': shape_location = (max_size, video_height - max_size) elif shape_location == 'bottom-right': shape_location = (video_width - max_size, video_height - max_size) else: raise ValueError(f"Invalid shape_location: {shape_location}") else: assert len(shape_location) == 2, "shape_location must be a tuple of (x, y) coordinates." shape_location = tuple(shape_location) # Convert shape_color from string to BGR for OpenCV color_map = {"red": (0, 0, 255), "green": (0, 255, 0), "blue": (255, 0, 0)} # Extendable shape_color_bgr = color_map.get(shape_color.lower(), (0, 0, 255)) # Default is red def add_shape_to_frame(get_frame, t): """ Function to add an expanding shape (circle, square, or triangle) to each frame based on time t. """ frame = get_frame(t) # Create a copy of the frame to modify frame = np.array(frame) # Calculate the current size of the shape based on the time elapsed current_size = int(min_size + (max_size - min_size) * (t / video_duration)) # Get the center for placing the shape center = shape_location # Draw the specified shape if shape == 'circle': cv2.circle(frame, center, current_size, shape_color_bgr, thickness) # Filled circle elif shape == 'square': top_left = (center[0] - current_size, center[1] - current_size) bottom_right = (center[0] + current_size, center[1] + current_size) cv2.rectangle(frame, top_left, bottom_right, shape_color_bgr, thickness) # Filled square elif shape == 'triangle': # Define triangle points p1 = (center[0], center[1] - current_size) # Top point p2 = (center[0] - current_size, center[1] + current_size) # Bottom-left point p3 = (center[0] + current_size, center[1] + current_size) # Bottom-right point pts = np.array([p1, p2, p3], np.int32) pts = pts.reshape((-1, 1, 2)) cv2.fillPoly(frame, [pts], shape_color_bgr) # Filled triangle return frame # Apply the expanding shape to each frame in the video video_with_shape = video_clip.fl(add_shape_to_frame) # Save the modified video with the expanding shape to the output path video_with_shape.write_videofile(save_path, codec="libx264", verbose=False, logger=None)