(args)
| 75 | random.seed(seed) |
| 76 | |
| 77 | def main(args): |
| 78 | device = torch.device(args.device) |
| 79 | |
| 80 | seed = args.seed |
| 81 | random_seed(seed) |
| 82 | |
| 83 | if args.model_size == 'base': |
| 84 | config = AutoConfig.from_pretrained('CLIP-B-16/config.json') |
| 85 | else: |
| 86 | raise NotImplementedError |
| 87 | |
| 88 | if args.model == 'anytouch': |
| 89 | |
| 90 | # tube_size is unused |
| 91 | model = TactileVideoMAE(args, config, args.num_frames, 1) # tube_size=1 is actually not used |
| 92 | |
| 93 | load_dir = args.load_path |
| 94 | ckpt = torch.load(load_dir, map_location='cpu') |
| 95 | model = load_model_from_multi_clip(ckpt, model) |
| 96 | print(load_dir) |
| 97 | else: |
| 98 | raise NotImplementedError(f'Model {args.model} not implemented!') |
| 99 | |
| 100 | model.to(device) |
| 101 | data, data_transformed = load_data(args) |
| 102 | Batch_size = 1 |
| 103 | data_transformed = data_transformed.unsqueeze(0).to(device) # B, T, C, H, W |
| 104 | print('Input Data shape:', data_transformed.shape) # T, C, H, W |
| 105 | |
| 106 | print('Visualization starting...') |
| 107 | vis_dir = args.output_dir |
| 108 | os.makedirs(vis_dir, exist_ok=True) |
| 109 | for i in range(data.shape[0]): |
| 110 | plt.imsave(vis_dir + f"/input_{i}.png", data[i].permute(1, 2, 0).numpy()) |
| 111 | plt.close() |
| 112 | print(f'Input images saved to {vis_dir}') |
| 113 | |
| 114 | print('Getting Sensor IDs') |
| 115 | print(f'Using sensor: {args.data_sensor}') |
| 116 | print(f'Sensor ID: {sensor_name_to_id[args.data_sensor]}') |
| 117 | print('Sensor ID mapping:', sensor_name_to_id) |
| 118 | sensor_id = sensor_name_to_id[args.data_sensor] |
| 119 | sensor_id_tensor = torch.ones((Batch_size,), dtype=torch.long, device=device) * sensor_id |
| 120 | print('Sensor ID tensor shape:', sensor_id_tensor.shape) |
| 121 | |
| 122 | print('Inference...') |
| 123 | model.eval() |
| 124 | with torch.no_grad(): |
| 125 | outputs = model(data_transformed, sensor_id_tensor, probe=True) |
| 126 | print('Model output feature shape (before projection):', outputs.shape) |
| 127 | ## Should be (1, 398, 768) for 4frames model. 398 = 1 (cls token) + 5 (sensor tokens) + 196 (patches) * 2 (time dim) |
| 128 | |
| 129 | outputs = model(data_transformed, sensor_id_tensor) |
| 130 | print('Model output feature shape (after projection):', outputs.shape) |
| 131 | ## Should be (1, 398, 512) for 4frames model. 398 = 1 (cls token) + 5 (sensor tokens) + 196 (patches) * 2 (time dim) |
| 132 | |
| 133 | cls_token = outputs[:, 0, :] # B, D |
| 134 | print('CLS token shape:', cls_token.shape) |
no test coverage detected