def get_test_loader(cfg, num_gpu, local_rank, stage, is_dist=True): # -------- get raw dataset interface -------- # normalize = transforms.Normalize(mean=cfg.INPUT.MEANS, std=cfg.INPUT.STDS) transform = transforms.Compose([transforms.ToTensor(), normalize]) attr = load_dataset(cfg.DATASET.NAME) if cfg.DATASET.NAME == 'COCO': Dataset = COCODataset elif cfg.DATASET.NAME == 'MPII': Dataset = MPIIDataset dataset = Dataset(attr, stage, transform) # -------- split dataset to gpus -------- # num_data = dataset.__len__() num_data_per_gpu = math.ceil(num_data / num_gpu) st = local_rank * num_data_per_gpu ed = min(num_data, st + num_data_per_gpu) indices = range(st, ed) subset = torch.utils.data.Subset(dataset, indices) # -------- make samplers -------- # sampler = torch.utils.data.sampler.SequentialSampler(subset) images_per_gpu = cfg.TEST.IMS_PER_GPU batch_sampler = torch.utils.data.sampler.BatchSampler(sampler, images_per_gpu, drop_last=False) # -------- make data_loader -------- # # 如何来组织这一个batch的数据 class BatchCollator(object): def __init__(self, size_divisible): self.size_divisible = size_divisible def __call__(self, batch): transposed_batch = list(zip(*batch)) images = torch.stack(transposed_batch[0], dim=0) scores = list(transposed_batch[1]) centers = list(transposed_batch[2]) scales = list(transposed_batch[3]) image_ids = list(transposed_batch[4]) return images, scores, centers, scales, image_ids data_loader = torch.utils.data.DataLoader( subset, num_workers=cfg.DATALOADER.NUM_WORKERS, batch_sampler=batch_sampler, collate_fn=BatchCollator(cfg.DATALOADER.SIZE_DIVISIBILITY), ) data_loader.ori_dataset = dataset return data_loader
def get_train_loader(cfg, num_gpu, is_dist=True, is_shuffle=True, start_iter=0): # -------- get raw dataset interface -------- # normalize = transforms.Normalize(mean=cfg.INPUT.MEANS, std=cfg.INPUT.STDS) transform = transforms.Compose([transforms.ToTensor(), normalize]) attr = load_dataset(cfg.DATASET.NAME) if cfg.DATASET.NAME == 'COCO': Dataset = COCODataset elif cfg.DATASET.NAME == 'MPII': Dataset = MPIIDataset dataset = Dataset(attr, 'train', transform) # -------- make samplers -------- # if is_dist: sampler = torch_samplers.DistributedSampler(dataset, shuffle=is_shuffle) elif is_shuffle: sampler = torch.utils.data.sampler.RandomSampler(dataset) else: sampler = torch.utils.data.sampler.SequentialSampler(dataset) images_per_gpu = cfg.SOLVER.IMS_PER_GPU # images_per_gpu = cfg.SOLVER.IMS_PER_BATCH // num_gpu aspect_grouping = [1] if cfg.DATALOADER.ASPECT_RATIO_GROUPING else [] if aspect_grouping: batch_sampler = torch_samplers.GroupedBatchSampler(sampler, dataset, aspect_grouping, images_per_gpu, drop_uneven=False) else: batch_sampler = torch.utils.data.sampler.BatchSampler(sampler, images_per_gpu, drop_last=False) batch_sampler = torch_samplers.IterationBasedBatchSampler( batch_sampler, cfg.SOLVER.MAX_ITER, start_iter) # -------- make data_loader -------- # class BatchCollator(object): def __init__(self, size_divisible): self.size_divisible = size_divisible def __call__(self, batch): transposed_batch = list(zip(*batch)) images = torch.stack(transposed_batch[0], dim=0) valids = torch.stack(transposed_batch[1], dim=0) labels = torch.stack(transposed_batch[2], dim=0) return images, valids, labels data_loader = torch.utils.data.DataLoader( dataset, num_workers=cfg.DATALOADER.NUM_WORKERS, batch_sampler=batch_sampler, collate_fn=BatchCollator(cfg.DATALOADER.SIZE_DIVISIBILITY), ) return data_loader
def visualize(self, img, joints, score=None): pairs = [[16, 14], [14, 12], [17, 15], [15, 13], [12, 13], [6, 12], [7, 13], [6, 7], [6, 8], [7, 9], [8, 10], [9, 11], [2, 3], [1, 2], [1, 3], [2, 4], [3, 5], [4, 6], [5, 7]] color = np.random.randint(0, 256, (self.keypoint_num, 3)).tolist() for i in range(self.keypoint_num): if joints[i, 0] > 0 and joints[i, 1] > 0: cv2.circle(img, tuple(joints[i, :2]), 2, tuple(color[i]), 2) if score: cv2.putText(img, score, (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (128, 255, 0), 2) def draw_line(img, p1, p2): c = (0, 0, 255) if p1[0] > 0 and p1[1] > 0 and p2[0] > 0 and p2[1] > 0: cv2.line(img, tuple(p1), tuple(p2), c, 2) for pair in pairs: draw_line(img, joints[pair[0] - 1], joints[pair[1] - 1]) return img if __name__ == '__main__': from dataset.attribute import load_dataset dataset = load_dataset('COCO') coco = COCODataset(dataset, 'val') print(coco.data_num)
class Config: # -------- Directoy Config -------- # USER = getpass.getuser() ROOT_DIR = os.environ['RSN_HOME'] OUTPUT_DIR = osp.join(ROOT_DIR, 'model_logs', USER, osp.split(osp.split(osp.realpath(__file__))[0])[1]) TEST_DIR = osp.join(OUTPUT_DIR, 'test_dir') TENSORBOARD_DIR = osp.join(OUTPUT_DIR, 'tb_dir') # -------- Data Config -------- # DATALOADER = edict() DATALOADER.NUM_WORKERS = 4 DATALOADER.ASPECT_RATIO_GROUPING = False DATALOADER.SIZE_DIVISIBILITY = 0 DATASET = edict() DATASET.NAME = 'COCO' dataset = load_dataset(DATASET.NAME) DATASET.KEYPOINT = dataset.KEYPOINT INPUT = edict() INPUT.NORMALIZE = True INPUT.MEANS = [0.406, 0.456, 0.485] # bgr INPUT.STDS = [0.225, 0.224, 0.229] # edict will automatcally convert tuple to list, so .. INPUT_SHAPE = dataset.INPUT_SHAPE OUTPUT_SHAPE = dataset.OUTPUT_SHAPE # -------- Model Config -------- # MODEL = edict() MODEL.BACKBONE = 'Res-50' MODEL.UPSAMPLE_CHANNEL_NUM = 256 MODEL.STAGE_NUM = 1 MODEL.OUTPUT_NUM = DATASET.KEYPOINT.NUM MODEL.DEVICE = 'cuda' MODEL.WEIGHT = None # -------- Training Config -------- # SOLVER = edict() SOLVER.BASE_LR = 5e-4 SOLVER.CHECKPOINT_PERIOD = 3200 SOLVER.GAMMA = 0.5 SOLVER.IMS_PER_GPU = 36 SOLVER.MAX_ITER = 128000 SOLVER.MOMENTUM = 0.9 SOLVER.OPTIMIZER = 'Adam' SOLVER.WARMUP_FACTOR = 0.1 SOLVER.WARMUP_ITERS = 2400 SOLVER.WARMUP_METHOD = 'linear' SOLVER.WEIGHT_DECAY = 1e-5 SOLVER.WEIGHT_DECAY_BIAS = 0 LOSS = edict() LOSS.OHKM = True LOSS.TOPK = 8 LOSS.COARSE_TO_FINE = True RUN_EFFICIENT = False # -------- Test Config -------- # TEST = dataset.TEST TEST.IMS_PER_GPU = 32
def visualize(self, img, joints, score=None): pairs = [[0, 1], [1, 2], [2, 6], [3, 4], [3, 6], [4, 5], [6, 7], [7, 8], [8, 9], [8, 12], [8, 13], [10, 11], [11, 12], [13, 14], [14, 15]] color = np.random.randint(0, 256, (self.keypoint_num, 3)).tolist() for i in range(self.keypoint_num): if joints[i, 0] > 0 and joints[i, 1] > 0: cv2.circle(img, tuple(joints[i, :2]), 2, tuple(color[i]), 2) if score: cv2.putText(img, score, (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (128, 255, 0), 2) def draw_line(img, p1, p2): c = (0, 0, 255) if p1[0] > 0 and p1[1] > 0 and p2[0] > 0 and p2[1] > 0: cv2.line(img, tuple(p1), tuple(p2), c, 2) for pair in pairs: draw_line(img, joints[pair[0] - 1], joints[pair[1] - 1]) return img if __name__ == '__main__': from dataset.attribute import load_dataset dataset = load_dataset('MPII') mpii = MPIIDataset(dataset, 'val') print(mpii.data_num)
ax.scatter(joints[:, 0], joints[:, 1]) for i in range(joints.shape[0]): ax.text(joints[i, 0], joints[i, 1], str(i), size=16, color='green') pass plt.savefig(outfile) # Manual testing if __name__ == '__main__': from dataset.attribute import load_dataset ### import matplotlib.pyplot as plt import matplotlib.image as mpimg from matplotlib.lines import Line2D ### dataset = load_dataset('H36M') h36m = H36MDataset(dataset, 'train') print(h36m.data_num) samples = [0, 70000, 433000, 190000, 600000] # samples = [0] for i, x in enumerate(samples): imgfile = 'train{}-{}.png'.format(x, i) data = h36m[x] print(data[0].shape) print(data[1].shape) print(data[3].shape) cv2.imwrite(imgfile, data[0]) __save_img(imgfile, data[3], 'markers-{}.png'.format(x))
for i in range(self.keypoint_num): if joints[i, 0] > 0 and joints[i, 1] > 0: cv.circle(img, tuple(joints[i, :2]), 2, tuple(color[i]), 2) if score: cv.putText(img, score, (50, 50), cv.FONT_HERSHEY_SIMPLEX, 1.2, (128, 255, 0), 2) def draw_line(img, p1, p2): c = (0, 0, 255) if p1[0] > 0 and p1[1] > 0 and p2[0] > 0 and p2[1] > 0: cv.line(img, tuple(p1), tuple(p2), c, 2) for pair in pairs: draw_line(img, joints[pair[0]], joints[pair[1]]) return img if __name__ == '__main__': from dataset.attribute import load_dataset dataset = load_dataset('MDJPE') mdj = MDJDataset(dataset, 'val') for i in range(len(mdj)): print(i) try: mdj[i] except: print(mdj[i]['img_path'])
[1, 2], [1, 3], [2, 4], [3, 5], [4, 6], [5, 7]] color = np.random.randint(0, 256, (self.keypoint_num, 3)).tolist() joints_array = np.ones((self.keypoint_num, 2), dtype=np.float32) for i in range(self.keypoint_num): joints_array[i, 0] = joints[i * 3] joints_array[i, 1] = joints[i * 3 + 1] # joints_array[i, 2] = joints[i * 3 + 2] for i in range(self.keypoint_num): if joints_array[i, 0] > 0 and joints_array[i, 1] > 0: cv2.circle(img, tuple(joints_array[i, :2]), 2, tuple(color[i]), 2) def draw_line(img, p1, p2): c = (0, 0, 255) if p1[0] > 0 and p1[1] > 0 and p2[0] > 0 and p2[1] > 0: cv2.line(img, tuple(p1), tuple(p2), c, 2) for pair in pairs: draw_line(img, joints_array[pair[0] - 1], joints_array[pair[1] - 1]) return img if __name__ == '__main__': normalize = transforms.Normalize(mean=cfg.INPUT.MEANS, std=cfg.INPUT.STDS) transform = transforms.Compose([transforms.ToTensor(), normalize]) attr = load_dataset(cfg.DATASET.NAME) dataset = PresentDataset(attr, cfg.INFO_PATH, transform)