Yuyangsb/cotracker
0
1# This Gradio demo code is from https://github.com/cvlab-kaist/locotrack/blob/main/demo/demo.py 2# We updated it to work with CoTracker3 models. We thank authors of LocoTrack3# for such an amazing Gradio demo.4 5import os6import sys7import uuid8 9import gradio as gr10import mediapy11import numpy as np12import cv213import matplotlib14import torch15import colorsys16import random17from typing import List, Optional, Sequence, Tuple18import spaces19import numpy as np20 21 22# Generate random colormaps for visualizing different points.23def get_colors(num_colors: int) -> List[Tuple[int, int, int]]:24 """Gets colormap for points."""25 colors = []26 for i in np.arange(0.0, 360.0, 360.0 / num_colors):27 hue = i / 360.028 lightness = (50 + np.random.rand() * 10) / 100.029 saturation = (90 + np.random.rand() * 10) / 100.030 color = colorsys.hls_to_rgb(hue, lightness, saturation)31 colors.append(32 (int(color[0] * 255), int(color[1] * 255), int(color[2] * 255))33 )34 random.shuffle(colors)35 return colors36 37def get_points_on_a_grid(38 size: int,39 extent: Tuple[float, ...],40 center: Optional[Tuple[float, ...]] = None,41 device: Optional[torch.device] = torch.device("cpu"),42):43 r"""Get a grid of points covering a rectangular region44 45 `get_points_on_a_grid(size, extent)` generates a :attr:`size` by46 :attr:`size` grid fo points distributed to cover a rectangular area47 specified by `extent`.48 49 The `extent` is a pair of integer :math:`(H,W)` specifying the height50 and width of the rectangle.51 52 Optionally, the :attr:`center` can be specified as a pair :math:`(c_y,c_x)`53 specifying the vertical and horizontal center coordinates. The center54 defaults to the middle of the extent.55 56 Points are distributed uniformly within the rectangle leaving a margin57 :math:`m=W/64` from the border.58 59 It returns a :math:`(1, \text{size} \times \text{size}, 2)` tensor of60 points :math:`P_{ij}=(x_i, y_i)` where61 62 .. math::63 P_{ij} = \left(64 c_x + m -\frac{W}{2} + \frac{W - 2m}{\text{size} - 1}\, j,~65 c_y + m -\frac{H}{2} + \frac{H - 2m}{\text{size} - 1}\, i66 \right)67 68 Points are returned in row-major order.69 70 Args:71 size (int): grid size.72 extent (tuple): height and with of the grid extent.73 center (tuple, optional): grid center.74 device (str, optional): Defaults to `"cpu"`.75 76 Returns:77 Tensor: grid.78 """79 if size == 1:80 return torch.tensor([extent[1] / 2, extent[0] / 2], device=device)[None, None]81 82 if center is None:83 center = [extent[0] / 2, extent[1] / 2]84 85 margin = extent[1] / 6486 range_y = (margin - extent[0] / 2 + center[0], extent[0] / 2 + center[0] - margin)87 range_x = (margin - extent[1] / 2 + center[1], extent[1] / 2 + center[1] - margin)88 grid_y, grid_x = torch.meshgrid(89 torch.linspace(*range_y, size, device=device),90 torch.linspace(*range_x, size, device=device),91 indexing="ij",92 )93 return torch.stack([grid_x, grid_y], dim=-1).reshape(1, -1, 2)94 95def paint_point_track(96 frames: np.ndarray,97 point_tracks: np.ndarray,98 visibles: np.ndarray,99 colormap: Optional[List[Tuple[int, int, int]]] = None,100) -> np.ndarray:101 """Converts a sequence of points to color code video.102 103 Args:104 frames: [num_frames, height, width, 3], np.uint8, [0, 255]105 point_tracks: [num_points, num_frames, 2], np.float32, [0, width / height]106 visibles: [num_points, num_frames], bool107 colormap: colormap for points, each point has a different RGB color.108 109 Returns:110 video: [num_frames, height, width, 3], np.uint8, [0, 255]111 """112 num_points, num_frames = point_tracks.shape[0:2]113 if colormap is None:114 colormap = get_colors(num_colors=num_points)115 height, width = frames.shape[1:3]116 dot_size_as_fraction_of_min_edge = 0.015117 radius = int(round(min(height, width) * dot_size_as_fraction_of_min_edge))118 diam = radius * 2 + 1119 quadratic_y = np.square(np.arange(diam)[:, np.newaxis] - radius - 1)120 quadratic_x = np.square(np.arange(diam)[np.newaxis, :] - radius - 1)121 icon = (quadratic_y + quadratic_x) - (radius**2) / 2.0122 sharpness = 0.15123 icon = np.clip(icon / (radius * 2 * sharpness), 0, 1)124 icon = 1 - icon[:, :, np.newaxis]125 icon1 = np.pad(icon, [(0, 1), (0, 1), (0, 0)])126 icon2 = np.pad(icon, [(1, 0), (0, 1), (0, 0)])127 icon3 = np.pad(icon, [(0, 1), (1, 0), (0, 0)])128 icon4 = np.pad(icon, [(1, 0), (1, 0), (0, 0)])129 130 video = frames.copy()131 for t in range(num_frames):132 # Pad so that points that extend outside the image frame don't crash us133 image = np.pad(134 video[t],135 [136 (radius + 1, radius + 1),137 (radius + 1, radius + 1),138 (0, 0),139 ],140 )141 for i in range(num_points):142 # The icon is centered at the center of a pixel, but the input coordinates143 # are raster coordinates. Therefore, to render a point at (1,1) (which144 # lies on the corner between four pixels), we need 1/4 of the icon placed145 # centered on the 0'th row, 0'th column, etc. We need to subtract146 # 0.5 to make the fractional position come out right.147 x, y = point_tracks[i, t, :] + 0.5148 x = min(max(x, 0.0), width)149 y = min(max(y, 0.0), height)150 151 if visibles[i, t]:152 x1, y1 = np.floor(x).astype(np.int32), np.floor(y).astype(np.int32)153 x2, y2 = x1 + 1, y1 + 1154 155 # bilinear interpolation156 patch = (157 icon1 * (x2 - x) * (y2 - y)158 + icon2 * (x2 - x) * (y - y1)159 + icon3 * (x - x1) * (y2 - y)160 + icon4 * (x - x1) * (y - y1)161 )162 x_ub = x1 + 2 * radius + 2163 y_ub = y1 + 2 * radius + 2164 image[y1:y_ub, x1:x_ub, :] = (1 - patch) * image[165 y1:y_ub, x1:x_ub, :166 ] + patch * np.array(colormap[i])[np.newaxis, np.newaxis, :]167 168 # Remove the pad169 video[t] = image[170 radius + 1 : -radius - 1, radius + 1 : -radius - 1171 ].astype(np.uint8)172 return video173 174 175PREVIEW_WIDTH = 768 # Width of the preview video176VIDEO_INPUT_RESO = (384, 512) # Resolution of the input video177POINT_SIZE = 4 # Size of the query point in the preview video178FRAME_LIMIT = 300 # Limit the number of frames to process179 180 181def get_point(frame_num, video_queried_preview, query_points, query_points_color, query_count, evt: gr.SelectData):182 print(f"You selected {(evt.index[0], evt.index[1], frame_num)}")183 184 current_frame = video_queried_preview[int(frame_num)]185 186 # Get the mouse click187 query_points[int(frame_num)].append((evt.index[0], evt.index[1], frame_num))188 189 # Choose the color for the point from matplotlib colormap190 color = matplotlib.colormaps.get_cmap("gist_rainbow")(query_count % 20 / 20)191 color = (int(color[0] * 255), int(color[1] * 255), int(color[2] * 255))192 # print(f"Color: {color}")193 query_points_color[int(frame_num)].append(color)194 195 # Draw the point on the frame196 x, y = evt.index197 current_frame_draw = cv2.circle(current_frame, (x, y), POINT_SIZE, color, -1)198 199 # Update the frame200 video_queried_preview[int(frame_num)] = current_frame_draw201 202 # Update the query count203 query_count += 1204 return (205 current_frame_draw, # Updated frame for preview206 video_queried_preview, # Updated preview video207 query_points, # Updated query points208 query_points_color, # Updated query points color209 query_count # Updated query count210 )211 212 213def undo_point(frame_num, video_preview, video_queried_preview, query_points, query_points_color, query_count):214 if len(query_points[int(frame_num)]) == 0:215 return (216 video_queried_preview[int(frame_num)],217 video_queried_preview,218 query_points,219 query_points_color,220 query_count221 )222 223 # Get the last point224 query_points[int(frame_num)].pop(-1)225 query_points_color[int(frame_num)].pop(-1)226 227 # Redraw the frame228 current_frame_draw = video_preview[int(frame_num)].copy()229 for point, color in zip(query_points[int(frame_num)], query_points_color[int(frame_num)]):230 x, y, _ = point231 current_frame_draw = cv2.circle(current_frame_draw, (x, y), POINT_SIZE, color, -1)232 233 # Update the query count234 query_count -= 1235 236 # Update the frame237 video_queried_preview[int(frame_num)] = current_frame_draw238 return (239 current_frame_draw, # Updated frame for preview240 video_queried_preview, # Updated preview video241 query_points, # Updated query points242 query_points_color, # Updated query points color243 query_count # Updated query count244 )245 246 247def clear_frame_fn(frame_num, video_preview, video_queried_preview, query_points, query_points_color, query_count):248 query_count -= len(query_points[int(frame_num)])249 250 query_points[int(frame_num)] = []251 query_points_color[int(frame_num)] = []252 253 video_queried_preview[int(frame_num)] = video_preview[int(frame_num)].copy()254 255 return (256 video_preview[int(frame_num)], # Set the preview frame to the original frame257 video_queried_preview, 258 query_points, # Cleared query points259 query_points_color, # Cleared query points color260 query_count # New query count261 )262 263 264 265def clear_all_fn(frame_num, video_preview):266 return (267 video_preview[int(frame_num)],268 video_preview.copy(),269 [[] for _ in range(len(video_preview))],270 [[] for _ in range(len(video_preview))],271 0272 )273 274 275def choose_frame(frame_num, video_preview_array):276 return video_preview_array[int(frame_num)]277 278 279def preprocess_video_input(video_path):280 video_arr = mediapy.read_video(video_path)281 video_fps = video_arr.metadata.fps282 num_frames = video_arr.shape[0]283 if num_frames > FRAME_LIMIT:284 gr.Warning(f"The video is too long. Only the first {FRAME_LIMIT} frames will be used.", duration=5)285 video_arr = video_arr[:FRAME_LIMIT]286 num_frames = FRAME_LIMIT287 288 # Resize to preview size for faster processing, width = PREVIEW_WIDTH289 height, width = video_arr.shape[1:3]290 new_height, new_width = int(PREVIEW_WIDTH * height / width), PREVIEW_WIDTH291 292 preview_video = mediapy.resize_video(video_arr, (new_height, new_width))293 input_video = mediapy.resize_video(video_arr, VIDEO_INPUT_RESO)294 295 preview_video = np.array(preview_video)296 input_video = np.array(input_video)297 298 interactive = True299 300 return (301 video_arr, # Original video302 preview_video, # Original preview video, resized for faster processing303 preview_video.copy(), # Copy of preview video for visualization304 input_video, # Resized video input for model305 # None, # video_feature, # Extracted feature306 video_fps, # Set the video FPS307 gr.update(open=False), # Close the video input drawer308 # tracking_mode, # Set the tracking mode309 preview_video[0], # Set the preview frame to the first frame310 gr.update(minimum=0, maximum=num_frames - 1, value=0, interactive=interactive), # Set slider interactive311 [[] for _ in range(num_frames)], # Set query_points to empty312 [[] for _ in range(num_frames)], # Set query_points_color to empty313 [[] for _ in range(num_frames)], 314 0, # Set query count to 0315 gr.update(interactive=interactive), # Make the buttons interactive316 gr.update(interactive=interactive),317 gr.update(interactive=interactive),318 gr.update(interactive=True),319 )320 321@spaces.GPU322def track(323 video_preview,324 video_input, 325 video_fps, 326 query_points, 327 query_points_color, 328 query_count, 329):330 tracking_mode = 'selected'331 if query_count == 0: 332 tracking_mode='grid'333 334 device = "cuda" if torch.cuda.is_available() else "cpu"335 dtype = torch.float if device == "cuda" else torch.float336 337 # Convert query points to tensor, normalize to input resolution338 if tracking_mode!='grid':339 query_points_tensor = []340 for frame_points in query_points:341 query_points_tensor.extend(frame_points)342 343 query_points_tensor = torch.tensor(query_points_tensor).float()344 query_points_tensor *= torch.tensor([345 VIDEO_INPUT_RESO[1], VIDEO_INPUT_RESO[0], 1346 ]) / torch.tensor([347 [video_preview.shape[2], video_preview.shape[1], 1]348 ])349 query_points_tensor = query_points_tensor[None].flip(-1).to(device, dtype) # xyt -> tyx350 query_points_tensor = query_points_tensor[:, :, [0, 2, 1]] # tyx -> txy351 352 video_input = torch.tensor(video_input).unsqueeze(0).to(device, dtype)353 354 model = torch.hub.load("facebookresearch/co-tracker", "cotracker3_online")355 model = model.to(device)356 357 video_input = video_input.permute(0, 1, 4, 2, 3)358 if tracking_mode=='grid':359 xy = get_points_on_a_grid(15, video_input.shape[3:], device=device)360 queries = torch.cat([torch.zeros_like(xy[:, :, :1]), xy], dim=2).to(device) #361 add_support_grid=False362 cmap = matplotlib.colormaps.get_cmap("gist_rainbow")363 query_points_color = [[]]364 query_count = queries.shape[1]365 for i in range(query_count):366 # Choose the color for the point from matplotlib colormap367 color = cmap(i / float(query_count))368 color = (int(color[0] * 255), int(color[1] * 255), int(color[2] * 255))369 query_points_color[0].append(color)370 371 else:372 queries = query_points_tensor373 add_support_grid=True374 375 model(video_chunk=video_input, is_first_step=True, grid_size=0, queries=queries, add_support_grid=add_support_grid)376 # 377 for ind in range(0, video_input.shape[1] - model.step, model.step):378 pred_tracks, pred_visibility = model(379 video_chunk=video_input[:, ind : ind + model.step * 2],380 grid_size=0, 381 queries=queries, 382 add_support_grid=add_support_grid383 ) # B T N 2, B T N 1384 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()385 pred_occ = pred_visibility[0].permute(1, 0).cpu().numpy()386 387 # make color array388 colors = []389 for frame_colors in query_points_color:390 colors.extend(frame_colors)391 colors = np.array(colors)392 393 painted_video = paint_point_track(video_preview,tracks,pred_occ,colors)394 395 # save video396 video_file_name = uuid.uuid4().hex + ".mp4"397 video_path = os.path.join(os.path.dirname(__file__), "tmp")398 video_file_path = os.path.join(video_path, video_file_name)399 os.makedirs(video_path, exist_ok=True)400 401 mediapy.write_video(video_file_path, painted_video, fps=video_fps)402 403 return video_file_path404 405 406with gr.Blocks() as demo:407 video = gr.State()408 video_queried_preview = gr.State()409 video_preview = gr.State()410 video_input = gr.State()411 video_fps = gr.State(24)412 413 query_points = gr.State([])414 query_points_color = gr.State([])415 is_tracked_query = gr.State([])416 query_count = gr.State(0)417 418 gr.Markdown("# 🎨 CoTracker3: Simpler and Better Point Tracking by Pseudo-Labelling Real Videos")419 gr.Markdown("<div style='text-align: left;'> \420 <p>Welcome to <a href='https://cotracker3.github.io/' target='_blank'>CoTracker</a>! This space demonstrates point (pixel) tracking in videos. \421 The model tracks points on a grid or points selected by you. </p> \422 <p> To get started, simply upload your <b>.mp4</b> video or click on one of the example videos to load them. The shorter the video, the faster the processing. We recommend submitting short videos of length <b>2-7 seconds</b>.</p> \423 <p> After you uploaded a video, please click \"Submit\" and then click \"Track\" for grid tracking or specify points you want to track before clicking. Enjoy the results! </p>\424 <p style='text-align: left'>For more details, check out our <a href='https://github.com/facebookresearch/co-tracker' target='_blank'>GitHub Repo</a> ⭐. We thank the authors of LocoTrack for their interactive demo.</p> \425 </div>"426 )427 428 429 gr.Markdown("## First step: upload your video or select an example video, and click submit.")430 with gr.Row():431 432 433 with gr.Accordion("Your video input", open=True) as video_in_drawer:434 video_in = gr.Video(label="Video Input", format="mp4")435 submit = gr.Button("Submit", scale=0)436 437 import os438 apple = os.path.join(os.path.dirname(__file__), "videos", "apple.mp4")439 bear = os.path.join(os.path.dirname(__file__), "videos", "bear.mp4")440 paragliding_launch = os.path.join(441 os.path.dirname(__file__), "videos", "paragliding-launch.mp4"442 )443 paragliding = os.path.join(os.path.dirname(__file__), "videos", "paragliding.mp4")444 cat = os.path.join(os.path.dirname(__file__), "videos", "cat.mp4")445 pillow = os.path.join(os.path.dirname(__file__), "videos", "pillow.mp4")446 teddy = os.path.join(os.path.dirname(__file__), "videos", "teddy.mp4")447 backpack = os.path.join(os.path.dirname(__file__), "videos", "backpack.mp4")448 449 450 gr.Examples(examples=[bear, apple, paragliding, paragliding_launch, cat, pillow, teddy, backpack], 451 inputs = [452 video_in453 ],454 )455 456 457 gr.Markdown("## Second step: Simply click \"Track\" to track a grid of points or select query points on the video before clicking")458 with gr.Row():459 with gr.Column():460 with gr.Row():461 query_frames = gr.Slider(462 minimum=0, maximum=100, value=0, step=1, label="Choose Frame", interactive=False)463 with gr.Row():464 undo = gr.Button("Undo", interactive=False)465 clear_frame = gr.Button("Clear Frame", interactive=False)466 clear_all = gr.Button("Clear All", interactive=False)467 468 with gr.Row():469 current_frame = gr.Image(470 label="Click to add query points", 471 type="numpy",472 interactive=False473 )474 475 with gr.Row():476 track_button = gr.Button("Track", interactive=False)477 478 with gr.Column():479 output_video = gr.Video(480 label="Output Video",481 interactive=False,482 autoplay=True,483 loop=True,484 )485 486 487 488 submit.click(489 fn = preprocess_video_input, 490 inputs = [video_in], 491 outputs = [492 video,493 video_preview,494 video_queried_preview,495 video_input,496 video_fps,497 video_in_drawer,498 current_frame,499 query_frames,500 query_points,501 query_points_color,502 is_tracked_query,503 query_count,504 undo,505 clear_frame,506 clear_all,507 track_button,508 ],509 queue = False510 )511 512 query_frames.change(513 fn = choose_frame,514 inputs = [query_frames, video_queried_preview],515 outputs = [516 current_frame,517 ],518 queue = False519 )520 521 current_frame.select(522 fn = get_point, 523 inputs = [524 query_frames,525 video_queried_preview,526 query_points,527 query_points_color,528 query_count,529 ], 530 outputs = [531 current_frame,532 video_queried_preview,533 query_points,534 query_points_color,535 query_count536 ], 537 queue = False538 )539 540 undo.click(541 fn = undo_point,542 inputs = [543 query_frames,544 video_preview,545 video_queried_preview,546 query_points,547 query_points_color,548 query_count549 ],550 outputs = [551 current_frame,552 video_queried_preview,553 query_points,554 query_points_color,555 query_count556 ],557 queue = False558 )559 560 clear_frame.click(561 fn = clear_frame_fn,562 inputs = [563 query_frames,564 video_preview,565 video_queried_preview,566 query_points,567 query_points_color,568 query_count569 ],570 outputs = [571 current_frame,572 video_queried_preview,573 query_points,574 query_points_color,575 query_count576 ],577 queue = False578 )579 580 clear_all.click(581 fn = clear_all_fn,582 inputs = [583 query_frames,584 video_preview,585 ],586 outputs = [587 current_frame,588 video_queried_preview,589 query_points,590 query_points_color,591 query_count592 ],593 queue = False594 )595 596 597 track_button.click(598 fn = track,599 inputs = [600 video_preview,601 video_input,602 video_fps,603 query_points,604 query_points_color,605 query_count,606 ],607 outputs = [608 output_video,609 ],610 queue = True,611 )612 613 614demo.launch(show_api=False, show_error=True, debug=False, share=False)