| 35 | |
| 36 | |
| 37 | def generate_latents_dinov2(args, batch, preprocessor, model, output_latents_path): |
| 38 | #encode img |
| 39 | rgb_img=batch["rgb_img"].cuda() |
| 40 | rgb_img=rgb_img[:,0:3,:,:] |
| 41 | |
| 42 | |
| 43 | rgb_input = preprocessor(rgb_img).to("cuda") |
| 44 | ret = model.forward_features(rgb_input) |
| 45 | patch_tok = ret["x_norm_patchtokens"].clone() |
| 46 | cls_tok = ret["x_norm_clstoken"].clone() |
| 47 | |
| 48 | # print("outputs",outputs) |
| 49 | |
| 50 | #only makes sense to write the last layer becuase it';s dinov2 and the other ones are not coarser representations |
| 51 | cls_token=cls_tok |
| 52 | # print("cls", cls_token.shape) |
| 53 | patch_embeddings = patch_tok |
| 54 | #reshape to [Batch_size, h, w, embedding] |
| 55 | batch_size, num_patches, hidden_size = patch_embeddings.shape |
| 56 | h = w = int(num_patches ** 0.5) # Assuming the number of patches is a perfect square (e.g., 14x14) |
| 57 | patch_embeddings_reshaped = patch_embeddings.reshape(batch_size, h, w, hidden_size) |
| 58 | patch_embeddings_reshaped=patch_embeddings_reshaped.permute(0,3,1,2).contiguous() #Make it bchw |
| 59 | #write last layer |
| 60 | out_path_final_latent=os.path.join(output_latents_path, "final_latent.pt") |
| 61 | torch.save(patch_embeddings_reshaped, out_path_final_latent) |
| 62 | #write cls token which is like an embedding for the whole image |
| 63 | out_path_cls_token=os.path.join(output_latents_path, "cls_token.pt") |
| 64 | #writing cls token of size |
| 65 | # print("cls_token",cls_token.shape) |
| 66 | torch.save(cls_token, out_path_cls_token) |
| 67 | |
| 68 | # feat_pca = img_2_pca(patch_embeddings_reshaped) |
| 69 | # torchvision.utils.save_image(feat_pca.squeeze(0), os.path.join(output_latents_path, "final_latent.png")) |
| 70 | |
| 71 | |
| 72 | |
| 73 | #write a file to signify that we are done with this folder |
| 74 | #start with x so that rsync copies it last if we copy to local |
| 75 | open( os.path.join(output_latents_path, "x_done.txt"), 'a').close() |
| 76 | |
| 77 | |
| 78 | def main(): |