def make_coco_transforms(image_set):
normalize = T.Compose([
T.ToTensor(),
T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
scales = [480, 512, 544, 576, 608, 640, 672, 704, 736, 768, 800]
if image_set == 'train':
return T.Compose([
T.RandomHorizontalFlip(),
T.RandomSelect(
T.RandomResize(scales, max_size=1333),
T.Compose([
T.RandomResize([400, 500, 600]),
T.RandomSizeCrop(384, 600),
T.RandomResize(scales, max_size=1333),
])),
normalize,
])
if image_set == 'val':
return T.Compose([
T.RandomResize([800], max_size=1333),
normalize,
])
raise ValueError(f'unknown {image_set}')
args.pin_memory = False
if not args.disable_cuda and torch.cuda.is_available():
args.device = torch.device('cuda')
args.pin_memory = True
return args
args = cli()
print("Loading dataset...")
# load train data
preprocess = transforms.Compose([
transforms.Normalize(),
transforms.RandomApply(transforms.HFlip(), 0.5),
transforms.RescaleRelative(),
transforms.Crop(args.square_edge),
transforms.CenterPad(args.square_edge),
])
class rtpose_lightning(pl.LightningModule):
def __init__(self, args, preprocess, target_transforms, model):
super(rtpose_lightning, self).__init__()
self.args = args
self.preprocess = preprocess
self.model = model
self.target_transforms = target_transforms
def forward(self, x):
last_ckpt_path = os.path.join(trainer.checkpoint_callback.filepath, last_ckpt_name)
state_dict = torch.load(last_ckpt_path, map_location=lambda storage, loc: storage)['state_dict']
os.remove(last_ckpt_path)
state_dict = {k.replace('model.',''):v for k,v in state_dict.items()}
rtpose_model.load_state_dict(state_dict)
return rtpose_model
args = parse_args()
update_config(cfg, args)
print("Loading dataset...")
# load train data
preprocess = transforms.Compose([
transforms.Normalize(),
transforms.RandomApply(transforms.HFlip(), 0.5),
transforms.RescaleRelative(scale_range=(cfg.DATASET.SCALE_MIN, cfg.DATASET.SCALE_MAX)),
transforms.Crop(cfg.DATASET.IMAGE_SIZE),
transforms.CenterPad(cfg.DATASET.IMAGE_SIZE),
])
# model
rtpose_vgg = get_model(trunk='vgg19')
# load pretrained
use_vgg(rtpose_vgg)
class rtpose_lightning(pl.LightningModule):
def __init__(self, preprocess, target_transforms, model, optimizer):
super(rtpose_lightning, self).__init__()
self.preprocess = preprocess
keypoints[part, 0] = joint_list[index, 0] + 0.5
keypoints[part, 1] = joint_list[index, 1] + 0.5
keypoints[part, 2] = 1.
outputs.append(keypoints)
return outputs
args = parse_args()
update_config(cfg, args)
print("Loading dataset...")
# load train data
preprocess = transforms.Compose([
transforms.Normalize(),
transforms.RandomRotate(max_rotate_degree=40),
transforms.RandomApply(transforms.HFlip(), 0.5),
transforms.RescaleRelative(),
transforms.Crop(cfg.DATASET.IMAGE_SIZE),
transforms.CenterPad(cfg.DATASET.IMAGE_SIZE),
])
def kp_connections(keypoints):
kp_lines = [
[keypoints.index('neck'), keypoints.index('right_hip')],
[keypoints.index('right_hip'), keypoints.index('right_knee')],
[keypoints.index('right_knee'), keypoints.index('right_ankle')],
[keypoints.index('neck'), keypoints.index('left_hip')],
[keypoints.index('left_hip'), keypoints.index('left_knee')],
[keypoints.index('left_knee'), keypoints.index('left_ankle')],
[keypoints.index('neck'), keypoints.index('right_shoulder')],
[keypoints.index('right_shoulder'), keypoints.index('right_elbow')],
if not args.disable_cuda and torch.cuda.is_available():
args.device = torch.device('cuda')
args.pin_memory = True
args.writer = tb.SummaryWriter(log_dir=args.log_dir) if args.log_dir else None
return args
args = cli()
print("Loading dataset...")
# load train data
preprocess = transforms.Compose([
transforms.Normalize(),
transforms.RandomApply(transforms.HFlip(), 0.5),
transforms.RescaleRelative(),
transforms.Crop(args.square_edge), # CHANGE args.square_edge (default 368) to larger for larger images!
transforms.CenterPad(args.square_edge),
])
train_loader, val_loader, train_data, val_data = train_factory(args, preprocess, target_transforms=None)
def build_names():
names = []
for j in range(1, 7):
for k in range(1, 3):
names.append('loss_stage%d_L%d' % (j, k))
return names
def get_loss(saved_for_loss, heat_temp, vec_temp):
