File size: 2,209 Bytes
18bb538 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
import os,sys
sys.path.append("..")
import numpy
from simple_dataset import Simple_InTheWild_dataset
import argparse
from torch.utils.data import DataLoader
import timm
import torch
import numpy as np
from util import misc
parser=argparse.ArgumentParser()
parser.add_argument("--data_dir",type=str,default="../example_process_data")
parser.add_argument('--world_size', default=1, type=int,
help='number of distributed processes')
parser.add_argument('--local_rank', default=-1, type=int)
parser.add_argument('--dist_on_itp', action='store_true')
parser.add_argument('--dist_url', default='env://',
help='url used to set up distributed training')
parser.add_argument('--scene_id',default="all",type=str)
args=parser.parse_args()
misc.init_distributed_mode(args)
dataset=Simple_InTheWild_dataset(dataset_dir=args.data_dir,scene_id=args.scene_id,n_px=224)
num_tasks = misc.get_world_size()
global_rank = misc.get_rank()
print(num_tasks,global_rank)
sampler = torch.utils.data.DistributedSampler(
dataset, num_replicas=num_tasks, rank=global_rank,
shuffle=False) # shuffle=True to reduce monitor bias
dataloader=DataLoader(
dataset,
sampler=sampler,
batch_size=10,
num_workers=4,
pin_memory=True,
drop_last=False
)
VIT_MODEL = 'vit_huge_patch14_224_clip_laion2b'
model=timm.create_model(VIT_MODEL, pretrained=True,pretrained_cfg_overlay=dict(file="./open_clip_pytorch_model.bin"))
model=model.eval().float().cuda()
for idx,data_batch in enumerate(dataloader):
if idx%10==0:
print("{}/{}".format(dataloader.__len__(),idx))
images = data_batch["images"].cuda().float()
model_id= data_batch["model_id"]
image_name=data_batch["image_name"]
scene_id=data_batch["scene_id"]
with torch.no_grad():
output_features=model.forward_features(images)
for j in range(output_features.shape[0]):
save_folder=os.path.join(args.data_dir,scene_id[j],"7_img_feature",model_id[j])
os.makedirs(save_folder,exist_ok=True)
save_path=os.path.join(save_folder,image_name[j]+".npz")
np.savez_compressed(save_path,img_features=output_features[j].detach().cpu().numpy().astype(np.float32))
|