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}')
# add args.device
args.device = torch.device('cpu')
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
