# this is built from https://huggingface.co/spaces/facebook/cotracker/blob/main/app.py # which was built from https://github.com/cvlab-kaist/locotrack/blob/main/demo/demo.py import os import sys import uuid from concurrent.futures import ThreadPoolExecutor import subprocess from nets.blocks import InputPadder import gradio as gr import mediapy import numpy as np import cv2 import matplotlib import torch import colorsys import random from typing import List, Optional, Sequence, Tuple import spaces import numpy as np import utils.py import utils.basic import utils.improc import PIL.Image # Generate random colormaps for visualizing different points. def get_colors(num_colors: int) -> List[Tuple[int, int, int]]: """Gets colormap for points.""" colors = [] for i in np.arange(0.0, 360.0, 360.0 / num_colors): hue = i / 360.0 lightness = (50 + np.random.rand() * 10) / 100.0 saturation = (90 + np.random.rand() * 10) / 100.0 color = colorsys.hls_to_rgb(hue, lightness, saturation) colors.append( (int(color[0] * 255), int(color[1] * 255), int(color[2] * 255)) ) random.shuffle(colors) return colors # def get_points_on_a_grid( # size: int, # extent: Tuple[float, ...], # center: Optional[Tuple[float, ...]] = None, # device: Optional[torch.device] = torch.device("cpu"), # ): # r"""Get a grid of points covering a rectangular region # `get_points_on_a_grid(size, extent)` generates a :attr:`size` by # :attr:`size` grid fo points distributed to cover a rectangular area # specified by `extent`. # The `extent` is a pair of integer :math:`(H,W)` specifying the height # and width of the rectangle. # Optionally, the :attr:`center` can be specified as a pair :math:`(c_y,c_x)` # specifying the vertical and horizontal center coordinates. The center # defaults to the middle of the extent. # Points are distributed uniformly within the rectangle leaving a margin # :math:`m=W/64` from the border. # It returns a :math:`(1, \text{size} \times \text{size}, 2)` tensor of # points :math:`P_{ij}=(x_i, y_i)` where # .. math:: # P_{ij} = \left( # c_x + m -\frac{W}{2} + \frac{W - 2m}{\text{size} - 1}\, j,~ # c_y + m -\frac{H}{2} + \frac{H - 2m}{\text{size} - 1}\, i # \right) # Points are returned in row-major order. # Args: # size (int): grid size. # extent (tuple): height and with of the grid extent. # center (tuple, optional): grid center. # device (str, optional): Defaults to `"cpu"`. # Returns: # Tensor: grid. # """ # if size == 1: # return torch.tensor([extent[1] / 2, extent[0] / 2], device=device)[None, None] # if center is None: # center = [extent[0] / 2, extent[1] / 2] # margin = extent[1] / 64 # range_y = (margin - extent[0] / 2 + center[0], extent[0] / 2 + center[0] - margin) # range_x = (margin - extent[1] / 2 + center[1], extent[1] / 2 + center[1] - margin) # grid_y, grid_x = torch.meshgrid( # torch.linspace(*range_y, size, device=device), # torch.linspace(*range_x, size, device=device), # indexing="ij", # ) # return torch.stack([grid_x, grid_y], dim=-1).reshape(1, -1, 2) @spaces.GPU def paint_point_track_gpu_scatter( frames: np.ndarray, point_tracks: np.ndarray, visibles: np.ndarray, colormap: Optional[List[Tuple[int, int, int]]] = None, rate: int = 1, show_bkg=True, # sharpness: float = 0.1, ) -> np.ndarray: print('starting vis') device = "cuda" if torch.cuda.is_available() else "cpu" frames_t = torch.from_numpy(frames).float().permute(0, 3, 1, 2).to(device) # [T,C,H,W] if show_bkg: frames_t = frames_t * 0.5 # darken, to see the point tracks better else: frames_t = frames_t * 0.0 # black out point_tracks_t = torch.from_numpy(point_tracks).to(device) # [P,T,2] visibles_t = torch.from_numpy(visibles).to(device) # [P,T] T, C, H, W = frames_t.shape P = point_tracks.shape[0] if colormap is None: colormap = get_colors(P) colors = torch.tensor(colormap, dtype=torch.float32, device=device) # [P,3] if rate==1: radius = 1 elif rate==2: radius = 1 elif rate== 4: radius = 2 elif rate== 8: radius = 4 else: radius = 6 # radius = max(1, int(np.sqrt(rate))) sharpness = 0.15 + 0.05 * np.log2(rate) D = radius * 2 + 1 y = torch.arange(D, device=device).float()[:, None] - radius x = torch.arange(D, device=device).float()[None, :] - radius dist2 = x**2 + y**2 icon = torch.clamp(1 - (dist2 - (radius**2) / 2.0) / (radius * 2 * sharpness), 0, 1) # [D,D] icon = icon.view(1, D, D) dx = torch.arange(-radius, radius + 1, device=device) dy = torch.arange(-radius, radius + 1, device=device) disp_y, disp_x = torch.meshgrid(dy, dx, indexing="ij") # [D,D] for t in range(T): mask = visibles_t[:, t] # [P] if mask.sum() == 0: continue xy = point_tracks_t[mask, t] + 0.5 # [N,2] xy[:, 0] = xy[:, 0].clamp(0, W - 1) xy[:, 1] = xy[:, 1].clamp(0, H - 1) colors_now = colors[mask] # [N,3] N = xy.shape[0] cx = xy[:, 0].long() # [N] cy = xy[:, 1].long() x_grid = cx[:, None, None] + disp_x # [N,D,D] y_grid = cy[:, None, None] + disp_y # [N,D,D] valid = (x_grid >= 0) & (x_grid < W) & (y_grid >= 0) & (y_grid < H) x_valid = x_grid[valid] # [K] y_valid = y_grid[valid] icon_weights = icon.expand(N, D, D)[valid] # [K] colors_valid = colors_now[:, :, None, None].expand(N, 3, D, D).permute(1, 0, 2, 3)[ :, valid ] # [3, K] idx_flat = (y_valid * W + x_valid).long() # [K] accum = torch.zeros_like(frames_t[t]) # [3, H, W] weight = torch.zeros(1, H * W, device=device) # [1, H*W] img_flat = accum.view(C, -1) # [3, H*W] weighted_colors = colors_valid * icon_weights # [3, K] img_flat.scatter_add_(1, idx_flat.unsqueeze(0).expand(C, -1), weighted_colors) weight.scatter_add_(1, idx_flat.unsqueeze(0), icon_weights.unsqueeze(0)) weight = weight.view(1, H, W) # accum = accum / (weight + 1e-6) # avoid division by 0 # frames_t[t] = torch.where(weight > 0, accum, frames_t[t]) # frames_t[t] = frames_t[t] * (1 - weight) + accum # alpha = weight.clamp(0, 1) # alpha = weight.clamp(0, 1) * 0.9 # transparency alpha = weight.clamp(0, 1) # transparency accum = accum / (weight + 1e-6) # [3, H, W] frames_t[t] = frames_t[t] * (1 - alpha) + accum * alpha # img_flat = frames_t[t].view(C, -1) # [3, H*W] # weighted_colors = colors_valid * icon_weights # [3, K] # img_flat.scatter_add_(1, idx_flat.unsqueeze(0).expand(C, -1), weighted_colors) print('done vis') return frames_t.clamp(0, 255).byte().permute(0, 2, 3, 1).cpu().numpy() def paint_point_track_gpu( frames: np.ndarray, point_tracks: np.ndarray, visibles: np.ndarray, colormap: Optional[List[Tuple[int, int, int]]] = None, radius: int = 2, sharpness: float = 0.15, ) -> np.ndarray: device = "cuda" if torch.cuda.is_available() else "cpu" # Setup frames_t = torch.from_numpy(frames).float().permute(0, 3, 1, 2).to(device) # [T,C,H,W] point_tracks_t = torch.from_numpy(point_tracks).to(device) # [P,T,2] visibles_t = torch.from_numpy(visibles).to(device) # [P,T] T, C, H, W = frames_t.shape P = point_tracks.shape[0] # Colors if colormap is None: colormap = get_colors(P) # or any fixed list of RGB colors = torch.tensor(colormap, dtype=torch.float32, device=device) # [P,3] # Icon kernel [K,K] D = radius * 2 + 1 y = torch.arange(D, device=device).float()[:, None] - radius - 1 x = torch.arange(D, device=device).float()[None, :] - radius - 1 dist2 = x**2 + y**2 icon = torch.clamp(1 - (dist2 - (radius**2) / 2.0) / (radius * 2 * sharpness), 0, 1) # [D,D] icon = icon.unsqueeze(0) # [1,D,D] for broadcasting # Create coordinate grids for t in range(T): image = frames_t[t] # Select visible points visible_mask = visibles_t[:, t] pt_xy = point_tracks_t[visible_mask, t] # [N,2] colors_t = colors[visible_mask] # [N,3] N = pt_xy.shape[0] if N == 0: continue # Integer centers pt_xy = pt_xy + 0.5 # correct center offset pt_xy[:, 0] = pt_xy[:, 0].clamp(0, W - 1) pt_xy[:, 1] = pt_xy[:, 1].clamp(0, H - 1) ix = pt_xy[:, 0].long() # [N] iy = pt_xy[:, 1].long() # Build grid of indices for patch around each point dx = torch.arange(-radius, radius + 1, device=device) dy = torch.arange(-radius, radius + 1, device=device) dx_grid, dy_grid = torch.meshgrid(dx, dy, indexing='ij') dx_flat = dx_grid.reshape(-1) dy_flat = dy_grid.reshape(-1) patch_x = ix[:, None] + dx_flat[None, :] # [N,K*K] patch_y = iy[:, None] + dy_flat[None, :] # [N,K*K] # Mask out-of-bounds valid = (patch_x >= 0) & (patch_x < W) & (patch_y >= 0) & (patch_y < H) flat_idx = (patch_y * W + patch_x).long() # [N,K*K] # Flatten icon and colors icon_flat = icon.view(1, -1) # [1, K*K] color_patches = colors_t[:, :, None] * icon_flat[:, None, :] # [N,3,K*K] # Flatten to write into 1D image img_flat = image.view(C, -1) # [3, H*W] for i in range(N): valid_mask = valid[i] idxs = flat_idx[i][valid_mask] vals = color_patches[i, :, valid_mask] # [3, valid_count] img_flat[:, idxs] += vals out_frames = frames_t.clamp(0, 255).byte().permute(0, 2, 3, 1).cpu().numpy() return out_frames def paint_point_track_parallel( frames: np.ndarray, point_tracks: np.ndarray, visibles: np.ndarray, colormap: Optional[List[Tuple[int, int, int]]] = None, max_workers: int = 8, ) -> np.ndarray: num_points, num_frames = point_tracks.shape[:2] if colormap is None: colormap = get_colors(num_colors=num_points) height, width = frames.shape[1:3] radius = 1 print('radius', radius) diam = radius * 2 + 1 # Precompute the icon and its bilinear components quadratic_y = np.square(np.arange(diam)[:, np.newaxis] - radius - 1) quadratic_x = np.square(np.arange(diam)[np.newaxis, :] - radius - 1) icon = (quadratic_y + quadratic_x) - (radius**2) / 2.0 sharpness = 0.15 icon = np.clip(icon / (radius * 2 * sharpness), 0, 1) icon = 1 - icon[:, :, np.newaxis] icon1 = np.pad(icon, [(0, 1), (0, 1), (0, 0)]) icon2 = np.pad(icon, [(1, 0), (0, 1), (0, 0)]) icon3 = np.pad(icon, [(0, 1), (1, 0), (0, 0)]) icon4 = np.pad(icon, [(1, 0), (1, 0), (0, 0)]) def draw_point(image, i, t): if not visibles[i, t]: return x, y = point_tracks[i, t, :] + 0.5 x = min(max(x, 0.0), width) y = min(max(y, 0.0), height) x1, y1 = np.floor(x).astype(np.int32), np.floor(y).astype(np.int32) x2, y2 = x1 + 1, y1 + 1 patch = ( icon1 * (x2 - x) * (y2 - y) + icon2 * (x2 - x) * (y - y1) + icon3 * (x - x1) * (y2 - y) + icon4 * (x - x1) * (y - y1) ) x_ub = x1 + 2 * radius + 2 y_ub = y1 + 2 * radius + 2 image[y1:y_ub, x1:x_ub, :] = (1 - patch) * image[y1:y_ub, x1:x_ub, :] + patch * np.array(colormap[i])[np.newaxis, np.newaxis, :] video = frames.copy() for t in range(num_frames): image = np.pad( video[t], [(radius + 1, radius + 1), (radius + 1, radius + 1), (0, 0)], ) with ThreadPoolExecutor(max_workers=max_workers) as executor: futures = [executor.submit(draw_point, image, i, t) for i in range(num_points)] _ = [f.result() for f in futures] # wait for all threads video[t] = image[radius + 1 : -radius - 1, radius + 1 : -radius - 1].astype(np.uint8) return video def paint_point_track( frames: np.ndarray, point_tracks: np.ndarray, visibles: np.ndarray, colormap: Optional[List[Tuple[int, int, int]]] = None, ) -> np.ndarray: """Converts a sequence of points to color code video. Args: frames: [num_frames, height, width, 3], np.uint8, [0, 255] point_tracks: [num_points, num_frames, 2], np.float32, [0, width / height] visibles: [num_points, num_frames], bool colormap: colormap for points, each point has a different RGB color. Returns: video: [num_frames, height, width, 3], np.uint8, [0, 255] """ num_points, num_frames = point_tracks.shape[0:2] if colormap is None: colormap = get_colors(num_colors=num_points) height, width = frames.shape[1:3] dot_size_as_fraction_of_min_edge = 0.015 # radius = int(round(min(height, width) * dot_size_as_fraction_of_min_edge)) radius = 2 # print('radius', radius) diam = radius * 2 + 1 quadratic_y = np.square(np.arange(diam)[:, np.newaxis] - radius - 1) quadratic_x = np.square(np.arange(diam)[np.newaxis, :] - radius - 1) icon = (quadratic_y + quadratic_x) - (radius**2) / 2.0 sharpness = 0.15 icon = np.clip(icon / (radius * 2 * sharpness), 0, 1) icon = 1 - icon[:, :, np.newaxis] icon1 = np.pad(icon, [(0, 1), (0, 1), (0, 0)]) icon2 = np.pad(icon, [(1, 0), (0, 1), (0, 0)]) icon3 = np.pad(icon, [(0, 1), (1, 0), (0, 0)]) icon4 = np.pad(icon, [(1, 0), (1, 0), (0, 0)]) video = frames.copy() for t in range(num_frames): # Pad so that points that extend outside the image frame don't crash us image = np.pad( video[t], [ (radius + 1, radius + 1), (radius + 1, radius + 1), (0, 0), ], ) for i in range(num_points): # The icon is centered at the center of a pixel, but the input coordinates # are raster coordinates. Therefore, to render a point at (1,1) (which # lies on the corner between four pixels), we need 1/4 of the icon placed # centered on the 0'th row, 0'th column, etc. We need to subtract # 0.5 to make the fractional position come out right. x, y = point_tracks[i, t, :] + 0.5 x = min(max(x, 0.0), width) y = min(max(y, 0.0), height) if visibles[i, t]: x1, y1 = np.floor(x).astype(np.int32), np.floor(y).astype(np.int32) x2, y2 = x1 + 1, y1 + 1 # bilinear interpolation patch = ( icon1 * (x2 - x) * (y2 - y) + icon2 * (x2 - x) * (y - y1) + icon3 * (x - x1) * (y2 - y) + icon4 * (x - x1) * (y - y1) ) x_ub = x1 + 2 * radius + 2 y_ub = y1 + 2 * radius + 2 image[y1:y_ub, x1:x_ub, :] = (1 - patch) * image[ y1:y_ub, x1:x_ub, : ] + patch * np.array(colormap[i])[np.newaxis, np.newaxis, :] # Remove the pad video[t] = image[ radius + 1 : -radius - 1, radius + 1 : -radius - 1 ].astype(np.uint8) return video PREVIEW_WIDTH = 1024 # Width of the preview video PREVIEW_HEIGHT = 1024 # VIDEO_INPUT_RESO = (384, 512) # Resolution of the input video POINT_SIZE = 1 # Size of the query point in the preview video FRAME_LIMIT = 600 # Limit the number of frames to process # def get_point(frame_num, video_queried_preview, query_points, query_points_color, query_count, evt: gr.SelectData): # print(f"You selected {(evt.index[0], evt.index[1], frame_num)}") # current_frame = video_queried_preview[int(frame_num)] # # Get the mouse click # query_points[int(frame_num)].append((evt.index[0], evt.index[1], frame_num)) # # Choose the color for the point from matplotlib colormap # color = matplotlib.colormaps.get_cmap("gist_rainbow")(query_count % 20 / 20) # color = (int(color[0] * 255), int(color[1] * 255), int(color[2] * 255)) # # print(f"Color: {color}") # query_points_color[int(frame_num)].append(color) # # Draw the point on the frame # x, y = evt.index # current_frame_draw = cv2.circle(current_frame, (x, y), POINT_SIZE, color, -1) # # Update the frame # video_queried_preview[int(frame_num)] = current_frame_draw # # Update the query count # query_count += 1 # return ( # current_frame_draw, # Updated frame for preview # video_queried_preview, # Updated preview video # query_points, # Updated query points # query_points_color, # Updated query points color # query_count # Updated query count # ) # def undo_point(frame_num, video_preview, video_queried_preview, query_points, query_points_color, query_count): # if len(query_points[int(frame_num)]) == 0: # return ( # video_queried_preview[int(frame_num)], # video_queried_preview, # query_points, # query_points_color, # query_count # ) # # Get the last point # query_points[int(frame_num)].pop(-1) # query_points_color[int(frame_num)].pop(-1) # # Redraw the frame # current_frame_draw = video_preview[int(frame_num)].copy() # for point, color in zip(query_points[int(frame_num)], query_points_color[int(frame_num)]): # x, y, _ = point # current_frame_draw = cv2.circle(current_frame_draw, (x, y), POINT_SIZE, color, -1) # # Update the query count # query_count -= 1 # # Update the frame # video_queried_preview[int(frame_num)] = current_frame_draw # return ( # current_frame_draw, # Updated frame for preview # video_queried_preview, # Updated preview video # query_points, # Updated query points # query_points_color, # Updated query points color # query_count # Updated query count # ) # def clear_frame_fn(frame_num, video_preview, video_queried_preview, query_points, query_points_color, query_count): # query_count -= len(query_points[int(frame_num)]) # query_points[int(frame_num)] = [] # query_points_color[int(frame_num)] = [] # video_queried_preview[int(frame_num)] = video_preview[int(frame_num)].copy() # return ( # video_preview[int(frame_num)], # Set the preview frame to the original frame # video_queried_preview, # query_points, # Cleared query points # query_points_color, # Cleared query points color # query_count # New query count # ) # def clear_all_fn(frame_num, video_preview): # return ( # video_preview[int(frame_num)], # video_preview.copy(), # [[] for _ in range(len(video_preview))], # [[] for _ in range(len(video_preview))], # 0 # ) def choose_frame(frame_num, video_preview_array): return video_preview_array[int(frame_num)] def choose_rate1(video_preview, video_fps, tracks, visibs): return choose_rate(1, video_preview, video_fps, tracks, visibs) def choose_rate2(video_preview, video_fps, tracks, visibs): return choose_rate(2, video_preview, video_fps, tracks, visibs) def choose_rate4(video_preview, video_fps, tracks, visibs): return choose_rate(4, video_preview, video_fps, tracks, visibs) def choose_rate8(video_preview, video_fps, tracks, visibs): return choose_rate(8, video_preview, video_fps, tracks, visibs) # def choose_rate16(video_preview, video_fps, tracks, visibs): # return choose_rate(16, video_preview, video_fps, tracks, visibs) def update_vis(rate, show_bkg, cmap, video_preview, query_frame, video_fps, tracks, visibs): print('rate', rate) print('cmap', cmap) print('video_preview', video_preview.shape) T, H, W,_ = video_preview.shape tracks_ = tracks.reshape(H,W,T,2)[::rate,::rate].reshape(-1,T,2) visibs_ = visibs.reshape(H,W,T)[::rate,::rate].reshape(-1,T) return paint_video(video_preview, query_frame, video_fps, tracks_, visibs_, rate=rate, show_bkg=show_bkg, cmap=cmap) # return video_preview_array[int(frame_num)] def preprocess_video_input(video_path): video_arr = mediapy.read_video(video_path) video_fps = video_arr.metadata.fps num_frames = video_arr.shape[0] if num_frames > FRAME_LIMIT: gr.Warning(f"The video is too long. Only the first {FRAME_LIMIT} frames will be used.", duration=5) video_arr = video_arr[:FRAME_LIMIT] num_frames = FRAME_LIMIT height, width = video_arr.shape[1:3] if height > width: new_height, new_width = PREVIEW_HEIGHT, int(PREVIEW_WIDTH * width / height) else: new_height, new_width = int(PREVIEW_WIDTH * height / width), PREVIEW_WIDTH if new_height*new_width > 768*1024: new_height = new_height*3//4 new_width = new_width*3//4 new_height, new_width = new_height//16 * 16, new_width//16 * 16 # make it divisible by 16, partly to satisfy ffmpeg preview_video = mediapy.resize_video(video_arr, (new_height, new_width)) # input_video = mediapy.resize_video(video_arr, VIDEO_INPUT_RESO) # input_video = video_arr input_video = preview_video preview_video = np.array(preview_video) input_video = np.array(input_video) return ( video_arr, # Original video preview_video, # Original preview video, resized for faster processing preview_video.copy(), # Copy of preview video for visualization input_video, # Resized video input for model video_fps, # Set the video FPS preview_video[0], # Set the preview frame to the first frame gr.update(minimum=0, maximum=num_frames - 1, value=0, interactive=True), # Set slider interactive gr.update(interactive=True), # make track button interactive # gr.update(interactive=True), # gr.update(interactive=True), # gr.update(interactive=True), # gr.update(interactive=True), ) def paint_video(video_preview, query_frame, video_fps, tracks, visibs, rate=1, show_bkg=True, cmap="gist_rainbow"): print('video_preview', video_preview.shape) print('tracks', tracks.shape) T, H, W, _ = video_preview.shape query_count = tracks.shape[0] print('cmap', cmap) print('query_frame', query_frame) if cmap=="bremm": # xy0 = tracks xy0 = tracks[:,query_frame] # N,2 # print('xyQ', xy0[:10]) # print('xy0', tracks[:10,0]) # print('xy1', tracks[:10,1]) colors = utils.improc.get_2d_colors(xy0, H, W) else: cmap_ = matplotlib.colormaps.get_cmap(cmap) query_points_color = [[]] for i in range(query_count): # Choose the color for the point from matplotlib colormap color = cmap_(i / float(query_count)) color = (int(color[0] * 255), int(color[1] * 255), int(color[2] * 255)) query_points_color[0].append(color) # make color array colors = [] for frame_colors in query_points_color: colors.extend(frame_colors) colors = np.array(colors) painted_video = paint_point_track_gpu_scatter(video_preview,tracks,visibs,colors,rate=rate,show_bkg=show_bkg)#=max(rate//2,1)) # save video video_file_name = uuid.uuid4().hex + ".mp4" video_path = os.path.join(os.path.dirname(__file__), "tmp") video_file_path = os.path.join(video_path, video_file_name) os.makedirs(video_path, exist_ok=True) if False: mediapy.write_video(video_file_path, painted_video, fps=video_fps) else: for ti in range(T): temp_out_f = '%s/%03d.jpg' % (video_path, ti) # temp_out_f = '%s/%03d.png' % (video_path, ti) im = PIL.Image.fromarray(painted_video[ti]) # im.save(temp_out_f, "PNG", subsampling=0, quality=80) im.save(temp_out_f) # print('saved', temp_out_f) # os.system('/usr/bin/ffmpeg -y -hide_banner -loglevel error -f image2 -framerate %d -pattern_type glob -i "%s/*.png" -c:v libx264 -crf 20 -pix_fmt yuv420p %s' % (video_fps, video_path, video_file_path)) os.system('/usr/bin/ffmpeg -y -hide_banner -loglevel error -f image2 -framerate %d -pattern_type glob -i "%s/*.jpg" -c:v libx264 -crf 20 -pix_fmt yuv420p %s' % (video_fps, video_path, video_file_path)) print('saved', video_file_path) for ti in range(T): # temp_out_f = '%s/%03d.png' % (video_path, ti) temp_out_f = '%s/%03d.jpg' % (video_path, ti) os.remove(temp_out_f) # print('deleted', temp_out_f) return video_file_path @spaces.GPU def track( video_preview, video_input, video_fps, query_frame, rate, show_bkg, cmap, ): # tracking_mode = 'selected' # if query_count == 0: # tracking_mode = 'grid' # print('query_frames', query_frames) # query_frame = int(query_frames[0]) # # query_frame = 0 device = "cuda" if torch.cuda.is_available() else "cpu" dtype = torch.float if device == "cuda" else torch.float print("0 torch.cuda.memory_allocated: %.1fGB"%(torch.cuda.memory_allocated(0)/1024/1024/1024)) # # Convert query points to tensor, normalize to input resolution # if tracking_mode!='grid': # query_points_tensor = [] # for frame_points in query_points: # query_points_tensor.extend(frame_points) # query_points_tensor = torch.tensor(query_points_tensor).float() # query_points_tensor *= torch.tensor([ # VIDEO_INPUT_RESO[1], VIDEO_INPUT_RESO[0], 1 # ]) / torch.tensor([ # [video_preview.shape[2], video_preview.shape[1], 1] # ]) # query_points_tensor = query_points_tensor[None].flip(-1).to(device, dtype) # xyt -> tyx # query_points_tensor = query_points_tensor[:, :, [0, 2, 1]] # tyx -> txy video_input = torch.tensor(video_input).unsqueeze(0).to(dtype) print("1 torch.cuda.memory_allocated: %.1fGB"%(torch.cuda.memory_allocated(0)/1024/1024/1024)) from nets.alltracker import Net model = Net(16) url = "https://huggingface.co/aharley/alltracker/resolve/main/alltracker.pth" state_dict = torch.hub.load_state_dict_from_url(url, map_location='cpu') model.load_state_dict(state_dict['model'], strict=True) print('loaded weights from', url) model = model.to(device) print("2 torch.cuda.memory_allocated: %.1fGB"%(torch.cuda.memory_allocated(0)/1024/1024/1024)) video_input = video_input.permute(0, 1, 4, 2, 3) print('video_input', video_input.shape) # model(video_input, iters=4, sw=None, is_training=False) # # model(video_chunk=video_input, is_first_step=True, grid_size=0, queries=queries, add_support_grid=add_support_grid) _, T, _, H, W = video_input.shape utils.basic.print_stats('video_input', video_input) print("3 torch.cuda.memory_allocated: %.1fGB"%(torch.cuda.memory_allocated(0)/1024/1024/1024)) grid_xy = utils.basic.gridcloud2d(1, H, W, norm=False, device='cpu:0').float() # 1,H*W,2 grid_xy = grid_xy.permute(0,2,1).reshape(1,1,2,H,W) # 1,1,2,H,W print("4 torch.cuda.memory_allocated: %.1fGB"%(torch.cuda.memory_allocated(0)/1024/1024/1024)) # if tracking_mode=='grid': # xy = get_points_on_a_grid(15, video_input.shape[3:], device=device) # queries = torch.cat([torch.zeros_like(xy[:, :, :1]), xy], dim=2).to(device) # # add_support_grid=False # cmap = matplotlib.colormaps.get_cmap("gist_rainbow") # query_points_color = [[]] # query_count = queries.shape[1] # for i in range(query_count): # # Choose the color for the point from matplotlib colormap # color = cmap(i / float(query_count)) # color = (int(color[0] * 255), int(color[1] * 255), int(color[2] * 255)) # query_points_color[0].append(color) # else: # queries = query_points_tensor # add_support_grid=True # query_frame = 0 torch.cuda.empty_cache() with torch.no_grad(): utils.basic.print_stats('video_input', video_input) if query_frame < T-1: flows_e, visconf_maps_e, _, _ = \ model(video_input[:, query_frame:], iters=4, sw=None, is_training=False) traj_maps_e = flows_e.cpu() + grid_xy # B,Tf,2,H,W visconf_maps_e = visconf_maps_e.cpu() else: traj_maps_e = torch.zeros((1,0,2,H,W), dtype=torch.float32) visconf_maps_e = torch.zeros((1,0,2,H,W), dtype=torch.float32) if query_frame > 0: backward_flows_e, backward_visconf_maps_e, _, _ = \ model(video_input[:, :query_frame+1].flip([1]), iters=4, sw=None, is_training=False) backward_traj_maps_e = backward_flows_e.cpu() + grid_xy # B,Tb,2,H,W, reversed backward_visconf_maps_e = backward_visconf_maps_e.cpu() backward_traj_maps_e = backward_traj_maps_e.flip([1]) # flip time backward_visconf_maps_e = backward_visconf_maps_e.flip([1]) # flip time if query_frame < T-1: backward_traj_maps_e = backward_traj_maps_e[:, :-1] # drop the overlapped frame backward_visconf_maps_e = backward_visconf_maps_e[:, :-1] # drop the overlapped frame traj_maps_e = torch.cat([backward_traj_maps_e, traj_maps_e], dim=1) # B,T,2,H,W visconf_maps_e = torch.cat([backward_visconf_maps_e, visconf_maps_e], dim=1) # B,T,2,H,W # if query_frame < T-1: # flows_e, visconf_maps_e, _, _ = \ # model.forward_sliding(video_input[:, query_frame:], iters=4, sw=None, is_training=False) # traj_maps_e = flows_e + grid_xy # B,Tf,2,H,W # print("5 torch.cuda.memory_allocated: %.1fGB"%(torch.cuda.memory_allocated(0)/1024/1024/1024)) # else: # traj_maps_e = torch.zeros((1,0,2,H,W), dtype=torch.float32) # visconf_maps_e = torch.zeros((1,0,2,H,W), dtype=torch.float32) # if query_frame > 0: # backward_flows_e, backward_visconf_maps_e, _, _ = \ # model.forward_sliding(video_input[:, :query_frame+1].flip([1]), iters=4, sw=None, is_training=False) # backward_traj_maps_e = backward_flows_e + grid_xy # B,Tb,2,H,W, reversed # backward_traj_maps_e = backward_traj_maps_e.flip([1]) # flip time # backward_visconf_maps_e = backward_visconf_maps_e.flip([1]) # flip time # if query_frame < T-1: # backward_traj_maps_e = backward_traj_maps_e[:, :-1] # drop the overlapped frame # backward_visconf_maps_e = backward_visconf_maps_e[:, :-1] # drop the overlapped frame # traj_maps_e = torch.cat([backward_traj_maps_e, traj_maps_e], dim=1) # B,T,2,H,W # visconf_maps_e = torch.cat([backward_visconf_maps_e, visconf_maps_e], dim=1) # B,T,2,H,W print("6 torch.cuda.memory_allocated: %.1fGB"%(torch.cuda.memory_allocated(0)/1024/1024/1024)) # for ind in range(0, video_input.shape[1] - model.step, model.step): # pred_tracks, pred_visibility = model( # video_chunk=video_input[:, ind : ind + model.step * 2], # grid_size=0, # queries=queries, # add_support_grid=add_support_grid # ) # B T N 2, B T N 1 # tracks = (pred_tracks * torch.tensor([video_preview.shape[2], video_preview.shape[1]]).to(device) / torch.tensor([VIDEO_INPUT_RESO[1], VIDEO_INPUT_RESO[0]]).to(device))[0].permute(1, 0, 2).cpu().numpy() # pred_occ = pred_visibility[0].permute(1, 0).cpu().numpy() # # make color array # colors = [] # for frame_colors in query_points_color: # colors.extend(frame_colors) # colors = np.array(colors) # traj_maps_e = traj_maps_e[:,:,:,::4,::4] # subsample # visconf_maps_e = visconf_maps_e[:,:,:,::4,::4] # subsample # traj_maps_e = traj_maps_e[:,:,:,::2,::2] # subsample # visconf_maps_e = visconf_maps_e[:,:,:,::2,::2] # subsample tracks = traj_maps_e.permute(0,3,4,1,2).reshape(-1,T,2).numpy() visibs = visconf_maps_e.permute(0,3,4,1,2).reshape(-1,T,2)[:,:,0].numpy() confs = visconf_maps_e.permute(0,3,4,1,2).reshape(-1,T,2)[:,:,0].numpy() # visibs = (visibs * confs) > 0.2 # N,T visibs = (confs) > 0.1 # N,T # visibs = (confs) > 0.1 # N,T # sc = (np.array([video_preview.shape[2], video_preview.shape[1]]) / np.array([VIDEO_INPUT_RESO[1], VIDEO_INPUT_RESO[0]])).reshape(1,1,2) # print('sc', sc) # tracks = tracks * sc return update_vis(rate, show_bkg, cmap, video_preview, query_frame, video_fps, tracks, visibs), tracks, visibs, gr.update(interactive=True), gr.update(interactive=True), gr.update(interactive=True) # gr.update(interactive=True), # gr.update(interactive=True), # gr.update(interactive=True), # gr.update(interactive=True), # gr.update(interactive=True)) # # query_count = tracks.shape[0] # query_count = tracks.shape[0] # cmap = matplotlib.colormaps.get_cmap("gist_rainbow") # query_points_color = [[]] # for i in range(query_count): # # Choose the color for the point from matplotlib colormap # color = cmap(i / float(query_count)) # color = (int(color[0] * 255), int(color[1] * 255), int(color[2] * 255)) # query_points_color[0].append(color) # # make color array # colors = [] # for frame_colors in query_points_color: # colors.extend(frame_colors) # colors = np.array(colors) # # visibs_ = visibs * 1.0 # # visibs_ = visibs_[:,1:] * visibs_[:,:-1] # # inds = np.sum(visibs_, axis=1) >= min(T//4,8) # # tracks = tracks[inds] # # visibs = visibs[inds] # # colors = colors[inds] # # painted_video = paint_point_track_parallel(video_preview,tracks,visibs,colors) # # painted_video = paint_point_track_gpu(video_preview,tracks,visibs,colors) # painted_video = paint_point_track_gpu_scatter(video_preview,tracks,visibs,colors) # print("7 torch.cuda.memory_allocated: %.1fGB"%(torch.cuda.memory_allocated(0)/1024/1024/1024)) # # save video # video_file_name = uuid.uuid4().hex + ".mp4" # video_path = os.path.join(os.path.dirname(__file__), "tmp") # video_file_path = os.path.join(video_path, video_file_name) # os.makedirs(video_path, exist_ok=True) # if False: # mediapy.write_video(video_file_path, painted_video, fps=video_fps) # else: # for ti in range(T): # temp_out_f = '%s/%03d.jpg' % (video_path, ti) # # temp_out_f = '%s/%03d.png' % (video_path, ti) # im = PIL.Image.fromarray(painted_video[ti]) # # im.save(temp_out_f, "PNG", subsampling=0, quality=80) # im.save(temp_out_f) # print('saved', temp_out_f) # # os.system('/usr/bin/ffmpeg -y -hide_banner -loglevel error -f image2 -framerate %d -pattern_type glob -i "%s/*.png" -c:v libx264 -crf 20 -pix_fmt yuv420p %s' % (video_fps, video_path, video_file_path)) # os.system('/usr/bin/ffmpeg -y -hide_banner -loglevel error -f image2 -framerate %d -pattern_type glob -i "%s/*.jpg" -c:v libx264 -crf 20 -pix_fmt yuv420p %s' % (video_fps, video_path, video_file_path)) # print('saved', video_file_path) # for ti in range(T): # # temp_out_f = '%s/%03d.png' % (video_path, ti) # temp_out_f = '%s/%03d.jpg' % (video_path, ti) # os.remove(temp_out_f) # print('deleted', temp_out_f) # # out_file = tempfile.NamedTemporaryFile(suffix="out.mp4", delete=False) # # subprocess.run(f"ffmpeg -y -loglevel quiet -stats -i {painted_video} -c:v libx264 {out_file.name}".split()) # return video_file_path with gr.Blocks() as demo: video = gr.State() video_queried_preview = gr.State() video_preview = gr.State() video_input = gr.State() video_fps = gr.State(24) # rate = gr.State([]) tracks = gr.State([]) visibs = gr.State([]) gr.Markdown("# ⚡ AllTracker: Efficient Dense Point Tracking at High Resolution") gr.Markdown("
\

This demo runs AllTracker to perform all-pixel tracking in a video of your choice.

\

To get started, simply upload an mp4, or select one of the example videos. The shorter the video, the faster the processing. We recommend submitting videos under 20 seconds long.

\

After picking a video, click \"Submit\" to load the frames into the app, and optionally choose a query frame (using the slider), and then click \"Track\".

\

For full info on how this works, check out our GitHub repo, or paper.

\

Initial code for this Gradio app came from LocoTrack and CoTracker -- big thanks to those authors!

\
" ) gr.Markdown("## Step 1: Select a video, and click \"Submit\".") with gr.Row(): with gr.Column(): with gr.Row(): video_in = gr.Video(label="Video input", format="mp4") with gr.Row(): submit = gr.Button("Submit") with gr.Column(): with gr.Row(): butterfly = os.path.join(os.path.dirname(__file__), "videos", "butterfly_800.mp4") monkey = os.path.join(os.path.dirname(__file__), "videos", "monkey_800.mp4") groundbox = os.path.join(os.path.dirname(__file__), "videos", "groundbox_800.mp4") apple = os.path.join(os.path.dirname(__file__), "videos", "apple.mp4") grasp_sponge_800 = os.path.join(os.path.dirname(__file__), "videos", "grasp_sponge_800.mp4") twist = os.path.join(os.path.dirname(__file__), "videos", "twist_800.mp4") parrot = os.path.join(os.path.dirname(__file__), "videos", "parrot_800.mp4") # dog = os.path.join(os.path.dirname(__file__), "videos", "dog.mp4") bear = os.path.join(os.path.dirname(__file__), "videos", "bear.mp4") paragliding_launch = os.path.join(os.path.dirname(__file__), "videos", "paragliding-launch.mp4") paragliding = os.path.join(os.path.dirname(__file__), "videos", "paragliding.mp4") cat = os.path.join(os.path.dirname(__file__), "videos", "cat.mp4") pillow = os.path.join(os.path.dirname(__file__), "videos", "pillow.mp4") teddy = os.path.join(os.path.dirname(__file__), "videos", "teddy.mp4") backpack = os.path.join(os.path.dirname(__file__), "videos", "backpack.mp4") gr.Examples(examples=[butterfly, parrot, groundbox, monkey, grasp_sponge_800, bear, apple, paragliding, paragliding_launch, cat, pillow, teddy, backpack, twist], inputs = [ video_in ], examples_per_page=20, ) # with gr.Column(): # gr.Markdown("Choose a video or upload one of your own.") gr.Markdown("## Step 2: Select a frame, and click \"Track\".") with gr.Row(): with gr.Column(): with gr.Row(): query_frame_slider = gr.Slider( minimum=0, maximum=100, value=0, step=1, label="Query frame", interactive=False) # with gr.Row(): # undo = gr.Button("Undo", interactive=False) # clear_frame = gr.Button("Clear Frame", interactive=False) # clear_all = gr.Button("Clear All", interactive=False) with gr.Row(): current_frame = gr.Image( # label="Click to add query points", label="Query frame", type="numpy", interactive=False ) with gr.Row(): track_button = gr.Button("Track", interactive=False) with gr.Column(): # with gr.Row(): # rate1_button = gr.Button("Subsampling", interactive=False) # rate2_button = gr.Button("Stride 2", interactive=False) # rate4_button = gr.Button("Rate 4", interactive=False) # rate8_button = gr.Button("Rate 8", interactive=False) # # rate16_button = gr.Button("Rate 16", interactive=False) with gr.Row(): # rate_slider = gr.Slider( # minimum=1, maximum=16, value=1, step=1, label="Choose subsampling rate", interactive=False) rate_radio = gr.Radio([1, 2, 4, 8, 16], value=2, label="Subsampling rate", interactive=False) with gr.Row(): cmap_radio = gr.Radio(["gist_rainbow", "rainbow", "jet", "turbo", "bremm"], value="gist_rainbow", label="Colormap", interactive=False) with gr.Row(): bkg_check = gr.Checkbox(value=False, label="Overlay tracks on video", interactive=False) with gr.Row(): output_video = gr.Video( label="Output video", interactive=False, autoplay=True, loop=True, ) submit.click( fn = preprocess_video_input, inputs = [video_in], outputs = [ video, video_preview, video_queried_preview, video_input, video_fps, current_frame, query_frame_slider, # undo, # clear_frame, # clear_all, track_button, ], queue = False ) query_frame_slider.change( fn = choose_frame, inputs = [query_frame_slider, video_queried_preview], outputs = [ current_frame, ], queue = False ) # current_frame.select( # fn = get_point, # inputs = [ # query_frames, # video_queried_preview, # query_points, # query_points_color, # query_count, # ], # outputs = [ # current_frame, # video_queried_preview, # query_points, # query_points_color, # query_count # ], # queue = False # ) # undo.click( # fn = undo_point, # inputs = [ # query_frames, # video_preview, # video_queried_preview, # query_points, # query_points_color, # query_count # ], # outputs = [ # current_frame, # video_queried_preview, # query_points, # query_points_color, # query_count # ], # queue = False # ) # clear_frame.click( # fn = clear_frame_fn, # inputs = [ # query_frames, # video_preview, # video_queried_preview, # query_points, # query_points_color, # query_count # ], # outputs = [ # current_frame, # video_queried_preview, # query_points, # query_points_color, # query_count # ], # queue = False # ) # clear_all.click( # fn = clear_all_fn, # inputs = [ # query_frames, # video_preview, # ], # outputs = [ # current_frame, # video_queried_preview, # query_points, # query_points_color, # query_count # ], # queue = False # ) # output_video = None track_button.click( fn = track, inputs = [ video_preview, video_input, video_fps, query_frame_slider, rate_radio, bkg_check, cmap_radio, ], outputs = [ output_video, tracks, visibs, rate_radio, bkg_check, cmap_radio, # rate1_button, # rate2_button, # rate4_button, # rate8_button, # rate16_button, ], queue = True, ) # rate_slider.change( # fn = choose_rate, # inputs = [rate_slider, video_preview, video_fps, tracks, visibs], # outputs = [ # output_video, # ], # queue = False # ) rate_radio.change( fn = update_vis, inputs = [rate_radio, bkg_check, cmap_radio, video_preview, query_frame_slider, video_fps, tracks, visibs], outputs = [ output_video, ], queue = False ) cmap_radio.change( fn = update_vis, inputs = [rate_radio, bkg_check, cmap_radio, video_preview, query_frame_slider, video_fps, tracks, visibs], outputs = [ output_video, ], queue = False ) bkg_check.change( fn = update_vis, inputs = [rate_radio, bkg_check, cmap_radio, video_preview, query_frame_slider, video_fps, tracks, visibs], outputs = [ output_video, ], queue = False ) # rate1_button.click( # fn = choose_rate1, # inputs = [video_preview, video_fps, tracks, visibs], # outputs = [output_video], # queue = False, # ) # rate2_button.click( # fn = choose_rate2, # inputs = [video_preview, video_fps, tracks, visibs], # outputs = [output_video], # queue = False, # ) # rate4_button.click( # fn = choose_rate4, # inputs = [video_preview, video_fps, tracks, visibs], # outputs = [output_video], # queue = False, # ) # rate8_button.click( # fn = choose_rate8, # inputs = [video_preview, video_fps, tracks, visibs], # outputs = [output_video], # queue = False, # ) # rate16_button.click( # fn = choose_rate16, # inputs = [video_preview, video_fps, tracks, visibs], # outputs = [output_video], # queue = False, # ) # demo.launch(show_api=False, show_error=True, debug=False, share=False) # demo.launch(show_api=False, show_error=True, debug=False, share=True) demo.launch(show_api=False, show_error=True, debug=False, share=False)