Update app.py

app.py

@@ -371,11 +371,11 @@ def track(
 
     else:
         queries = query_points_tensor
-        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)  #
+        # 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)  #
         num_tracks = queries.shape[1]
-        queries = torch.cat([queries,queries__],dim=1)
-        add_support_grid=True
+        # queries = torch.cat([queries,queries__],dim=1)
+        # add_support_grid=True
 
     model(video_chunk=video_input, is_first_step=True, grid_size=0, queries=queries, add_support_grid=add_support_grid)
     # 
@@ -387,7 +387,7 @@ def track(
             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()
+    pred_occ = torch.ones_like(pred_visibility[0]).permute(1, 0).cpu().numpy()
 
     # make color array
     colors = []
@@ -395,16 +395,16 @@ def track(
         colors.extend(frame_colors)
     colors = np.array(colors)
 
-    pred_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))
+    # pred_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))
     
-    vis = Visualizer(save_dir="./saved_videos", pad_value=0, linewidth=1, tracks_leave_trace=0)
+    # vis = Visualizer(save_dir="./saved_videos", pad_value=0, linewidth=1, tracks_leave_trace=0)
     # segm_mask = torch.zeros(queries.shape[1])
     # segm_mask[:num_tracks] = 1
     # print('segm_mask',segm_mask.shape, segm_mask)
     # segm_mask=segm_mask, 
-    painted_video = vis.visualize(torch.tensor(video_preview).permute(0, 3, 1, 2)[None].to(pred_tracks.device), pred_tracks, pred_visibility, save_video=False)[0].permute(0, 2, 3, 1).cpu().numpy()
+    # painted_video = vis.visualize(torch.tensor(video_preview).permute(0, 3, 1, 2)[None].to(pred_tracks.device), pred_tracks, pred_visibility, save_video=False)[0].permute(0, 2, 3, 1).cpu().numpy()
     
-    # painted_video = paint_point_track(video_preview,tracks,pred_occ,colors)
+    painted_video = paint_point_track(video_preview,tracks,pred_occ,colors)
 
     # save video
     video_file_name = uuid.uuid4().hex + ".mp4"