add a demo code
This commit is contained in:
		
							
								
								
									
										47
									
								
								demo1.py
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										47
									
								
								demo1.py
									
									
									
									
									
										Normal file
									
								
							| @@ -0,0 +1,47 @@ | ||||
| import os | ||||
| import torch | ||||
|  | ||||
| from base64 import b64encode | ||||
| from cotracker.utils.visualizer import Visualizer, read_video_from_path | ||||
|  | ||||
|  | ||||
| video = read_video_from_path('./assets/F1_shorts.mp4') | ||||
| video = torch.from_numpy(video).permute(0, 3, 1, 2)[None].float() | ||||
|  | ||||
| from cotracker.predictor import CoTrackerPredictor | ||||
|  | ||||
| model = CoTrackerPredictor( | ||||
|     checkpoint=os.path.join( | ||||
|         './checkpoints/cotracker2.pth' | ||||
|     ) | ||||
| ) | ||||
|  | ||||
| # pred_tracks, pred_visibility = model(video, grid_size=30) | ||||
|  | ||||
| # vis = Visualizer(save_dir='./videos', pad_value=100) | ||||
| # vis.visualize(video=video, tracks=pred_tracks, visibility=pred_visibility, filename='teaser'); | ||||
|  | ||||
| grid_query_frame=20 | ||||
|  | ||||
| import torch.nn.functional as F | ||||
| video_interp = F.interpolate(video[0], [200, 360], mode="bilinear")[None] | ||||
|  | ||||
| import time | ||||
| start_time = time.time() | ||||
| pred_tracks, pred_visibility = model(video_interp,  | ||||
|                                      grid_query_frame=grid_query_frame, backward_tracking=True) | ||||
| end_time = time.time()  | ||||
|  | ||||
| print("Time taken: {:.2f} seconds".format(end_time - start_time)) | ||||
|  | ||||
| vis = Visualizer( | ||||
|     save_dir='./videos', | ||||
|     pad_value=20, | ||||
|     linewidth=1, | ||||
|     mode='optical_flow' | ||||
| ) | ||||
| vis.visualize( | ||||
|     video=video, | ||||
|     tracks=pred_tracks, | ||||
|     visibility=pred_visibility, | ||||
|     filename='dense'); | ||||
		Reference in New Issue
	
	Block a user