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)
