naver-ai / relabel_imagenet

Apache License 2.0
396 stars 39 forks source link

Relabel maps generating code #10

Open rtaori opened 2 years ago

rtaori commented 2 years ago

Hi,

This is exciting work! Would you be able to please provide the code that you used to generate the relabel imagenet maps? I would like to apply this technique to my own dataset and the relabeling code would be super helpful. Even if it's not polished, it would be a great help :)

Thanks,

hellbell commented 2 years ago

@rtaori Unfortunately, we don't have a plan to update our codebase now.
I can share the label map extraction code snippet, but I'm not sure whether this is the right and valid code. I hope you can understand how the label maps were extracted through this code, but if you want fully-working codes, please wait for our next codebase update.

normalize = torchvision.transforms.Normalize(
    mean=[0.485, 0.456, 0.406],
    std=[0.229, 0.224, 0.225]
)
train_root = os.path.join(data_path, 'train')
train_dataset = ImageFolderWithPath(
    root=train_root,
    transform=torchvision.transforms.Compose([
        torchvision.transforms.Resize((machine_annotator_image_size,
                                       machine_annotator_image_size)),
        torchvision.transforms.ToTensor(),
        normalize,
    ]))

train_sampler = torch.utils.data.distributed.DistributedSampler(
    train_dataset)

train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_sizebatch_size,
    shuffle=False,
    num_workers=machine_annotator_num_workers,
    pin_memory=True,
    sampler=train_sampler
)

score_maps = {}

print('Train loader size: {}'.format(len(train_loader)))
for iteration, batch in enumerate(train_loader):
    start_time = time.time()

    input, target, img_names = batch
    target = target.cuda()

    with torch.no_grad():
        # get score_map by removing GlobalAvergaPool. (See Appendix)
        output_map = machine_annotator_model(input, score_map=True) 

    output_topk = output_map.topk(k=machine_annotator_topk, dim=1)

    for idx, (t0, t1) in enumerate(zip(output_topk[0], output_topk[1])):
        img_name = img_names[idx].split('/')[-1].split('.')[0]
        t = torch.cat([t0.unsqueeze(0), t1.unsqueeze(0).float()],
                      dim=0)

        score_maps[img_name] = t.cpu()

score_map_path = '{}_{}_sz{}_top{}_score_maps.pt'.format(
    data.dataset,
    machine_annotator_arch,
    machine_annotator_image_size,
    machine_annotator_topk)