cyang-cityu / metacorrection Goto Github PK

# self.img_ids = [i_id.strip() for i_id in open(list_path)]
folder_to_find_desired_images = str(Path(os.path.join(root, f'leftImg8bit/{set}/')))
self.img_ids = [str(path)[len(folder_to_find_desired_images)+1:] for path in Path(folder_to_find_desired_images).rglob('*.png')]

# In train_meta.py
def generate_pseudo(cfg, batch_size=4):
    if not os.path.exists(cfg['data']['target']['pseudo_label_folder']):
        os.makedirs(cfg['data']['target']['pseudo_label_folder'])
    device = torch.device("cuda")
    model = DeeplabMulti(num_classes=19).cuda()
    model.load_state_dict(torch.load(os.path.join(cfg['model']['pretrained_folder'], 'GTA5_best.pth')))
    model.eval()

    target_train_loader = data.DataLoader(cityscapesDataSet(cfg['data']['target']['rootpath'], '', crop_size=(1024, 512), mean=IMG_MEAN, scale=False, mirror=False, set='train'),
                                    batch_size=batch_size, shuffle=False, pin_memory=True)

    interp = nn.Upsample(size=(1024, 2048), mode='bilinear', align_corners=True)

    with torch.no_grad():
        pbar = tqdm(enumerate(target_train_loader), total=len(target_train_loader.dataset)//target_train_loader.batch_size, desc='Generating pseudo-labels...')
        for index, batch in pbar:
            image, _, names = batch
            image = image.to(device)
            output1, output2 = model(image)

            output = interp(0.45 * output2 + 0.55 * output1).cpu().data.numpy()
            for i, out in enumerate(output):
                out = out.transpose(1,2,0)
                out = np.asarray(np.argmax(out, axis=2), dtype=np.uint8)
                out = Image.fromarray(out)

                name = names[i].split('/')[-1]
                out.save('%s/%s' % (cfg['data']['target']['pseudo_label_folder'], name))

def main():

  # ...

  generate_pseudo(cfg)
  targetloader = data.DataLoader(cityscapesPseudo(args.data_dir_target, args.data_list_target, cfg['data']['target']['pseudo_label_folder'],
  max_iters=args.num_steps * args.iter_size * args.batch_size,
  crop_size=input_size_target,
  scale=False, mirror=args.random_mirror, mean=IMG_MEAN, split='train'),
  batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers,
  pin_memory=True)

  # ...

class cityscapesPseudo(data.Dataset):
    def __init__(self, root, list_path, pseudo_path, max_iters=None, crop_size=(321, 321), mean=(128, 128, 128), scale=True, mirror=True, ignore_label=255, split='train'):
        self.root = root
        self.list_path = list_path
        self.crop_size = crop_size
        self.scale = scale
        self.ignore_label = ignore_label
        self.mean = mean
        self.is_mirror = mirror
        # self.mean_bgr = np.array([104.00698793, 116.66876762, 122.67891434])
        # self.img_ids = [i_id.strip().split() for i_id in open(list_path)]
        folder_to_find_desired_images = str(Path(os.path.join(root, f'leftImg8bit/{split}/')))
        self.img_ids = [str(path)[len(folder_to_find_desired_images)+1:] for path in Path(folder_to_find_desired_images).rglob('*.png')]
        if not max_iters==None:
            self.img_ids = self.img_ids * int(np.ceil(float(max_iters) / len(self.img_ids)))

        self.files = []
        self.set = set
        # for split in ["train", "trainval", "val"]:
        for name in self.img_ids:
            img_file = osp.join(self.root, "leftImg8bit/%s/%s" % (split, name))
            label_file = osp.join(pseudo_path, osp.basename(name))
            self.files.append({
                "img": img_file,
                "label": label_file,
                "name": name
            })
        # print(self.files)

    def __len__(self):
        return len(self.files)

    def __getitem__(self, index):
        datafiles = self.files[index]

        image = Image.open(datafiles["img"]).convert('RGB')
        label = Image.open(datafiles["label"])
        name = datafiles["name"]

        # resize
        image = image.resize(self.crop_size, Image.BICUBIC)
        label = label.resize(self.crop_size, Image.NEAREST)

        #image, label = randomCrop(image, label, self.crop_size[0], self.crop_size[1])

        image = np.asarray(image, np.float32)
        label = np.asarray(label, np.float32)

        if self.is_mirror:
            flip = np.random.choice(2) * 2 - 1
            image = image[:, :, ::flip]
            label = label[:, ::flip]

        size = image.shape
        image = image[:, :, ::-1]  # change to BGR
        image -= self.mean
        image = image.transpose((2, 0, 1))

        return image.copy(), label.copy(), np.array(size), name