def __init__(self, train_loader, val_loader, args, njoints, device, idx):
self.train_loader = train_loader
self.val_loader = val_loader
self.train_num_batches = len(self.train_loader)
self.val_num_batches = len(self.val_loader)
self.batch_size = args.train_batch
self.num_classes = njoints
self.device = device
self.idx = idx
self.epochs = args.epochs
self.schedule = args.schedule
self.gamma = args.gamma
self.sigma_decay = args.sigma_decay
self.checkpoint = args.checkpoint
self.snapshot = args.snapshot
self.num_stacks = args.stacks
self.debug = args.debug
self.flip = args.flip
print("==> creating model '{}', stacks={}, blocks={}".format(
args.arch, args.stacks, args.blocks))
self.netG, self.netsD, self.start_epoch, self.best_acc, self.logger, self.optimizerG, self.optimizersD, self.domainD = \
self.build_models(num_stacks=args.stacks,
num_blocks=args.blocks,
num_classes=njoints,
resnet_layers=args.resnet_layers,
device=device,
args=args)
self.real_labels, self.fake_labels = self.prepare_labels()
self.criterion = losses.JointsMSELoss().to(device)
def forward_pass(args):
njoints = datasets.__dict__[args.dataset].njoints
print("==> creating model '{}', stacks={}, blocks={}".format(
args.arch, args.stacks, args.blocks))
model = models.__dict__[args.arch](num_stacks=args.stacks,
num_blocks=args.blocks,
num_classes=njoints,
resnet_layers=args.resnet_layers)
model = torch.nn.DataParallel(model).to(device)
# define loss function (criterion) and optimizer
criterion = losses.JointsMSELoss().to(device)
if args.solver == 'rms':
optimizer = torch.optim.RMSprop(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.solver == 'adam':
optimizer = torch.optim.Adam(
model.parameters(),
lr=args.lr,
)
else:
print('Unknown solver: {}'.format(args.solver))
assert False
lr = args.lr
def __init__(self):
self.device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
cudnn.benchmark = True
self.img_path = ''
self.dataset = 'mpii'
self.image_path = ''
self.inp_res = 256
self.out_res = 64
self.arch = 'hg'
self.stacks = 2
self.blocks = 1
self.features = 256
self.resnet_layers = 50
self.solver = 'rms'
self.workers = 1
self.epochs = 100
self.test_batch = 1
self.train_batch = 1
self.lr = 2.5e-4
self.momentum = 0
self.weight_decay = 0
self.gamma = 0.1
self.sigma = 1.0
self.scale_factor = 0.25
self.rot_factor = 1
self.sigma_decay = 0
#self.checkpoint=''
self.resume = ''
self.njoints = 24
self.model = models.hg(num_stacks=self.stacks,
num_blocks=self.blocks,
num_classes=self.njoints,
resnet_layers=self.resnet_layers)
self.model = torch.nn.DataParallel(self.model).to(self.device)
self.criterion = losses.JointsMSELoss().to(self.device)
self.optimizer = torch.optim.RMSprop(self.model.parameters(),
lr=self.lr,
momentum=self.momentum,
weight_decay=self.weight_decay)
# TODO changing hardcoded model path to parser
self.checkpoint = torch.load(
"/home/rbccps2080ti/projects/link_speed_estimation/ta_darknet/hourglass/checkpoint/mpii/hg_updated_21/checkpoint.pth.tar"
)
self.start_epoch = self.checkpoint['epoch']
self.model.load_state_dict(self.checkpoint['state_dict'])
self.model.eval()
def main(args):
global best_acc
global idx1
global idx2
# idx is the index of joints used to compute accuracy for dataset2
idx1 = range(1, 19)
idx2 = range(1, 19) # horse
# create model
njoints = datasets.__dict__[args.dataset].njoints
print("==> creating model '{}'".format(args.arch))
model = models.__dict__[args.arch](num_classes=njoints,
resnet_layers=args.resnet_layers,
pretrained=None,
dual_branch=True)
model = torch.nn.DataParallel(model).to(device)
# define loss function (criterion) and optimizer
criterion = losses.JointsMSELoss().to(device)
# optionally resume from a checkpoint
if args.resume:
if isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_acc = checkpoint['best_acc']
model.load_state_dict(checkpoint['state_dict_ema'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
else:
raise Exception('please provide a checkpoint')
val_dataset = datasets.__dict__[args.dataset](is_train=False,
is_aug=False,
**vars(args))
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
_, acc, predictions = validate(val_loader, model, criterion, njoints, args,
args.flip, args.test_batch)
return
def main(args):
global best_acc
global idx
# idx is the index of joints used to compute accuracy
if args.dataset in ['mpii', 'lsp']:
idx = [1, 2, 3, 4, 5, 6, 11, 12, 15, 16]
elif args.dataset == 'coco':
idx = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17]
else:
print("Unknown dataset: {}".format(args.dataset))
assert False
# create checkpoint dir
if not isdir(args.checkpoint):
mkdir_p(args.checkpoint)
# create model
njoints = datasets.__dict__[args.dataset].njoints
print("==> creating model '{}', stacks={}, blocks={}".format(
args.arch, args.stacks, args.blocks))
model = models.__dict__[args.arch](num_stacks=args.stacks,
num_blocks=args.blocks,
num_classes=njoints,
resnet_layers=args.resnet_layers)
model = torch.nn.DataParallel(model).to(device)
# define loss function (criterion) and optimizer
criterion = losses.JointsMSELoss().to(device)
if args.solver == 'rms':
optimizer = torch.optim.RMSprop(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.solver == 'adam':
optimizer = torch.optim.Adam(
model.parameters(),
lr=args.lr,
)
else:
print('Unknown solver: {}'.format(args.solver))
assert False
# optionally resume from a checkpoint
title = args.dataset + ' ' + args.arch
if args.resume:
if isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_acc = checkpoint['best_acc']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
logger = Logger(join(args.checkpoint, 'log.txt'),
title=title,
resume=True)
else:
print("=> no checkpoint found at '{}'".format(args.resume))
else:
logger = Logger(join(args.checkpoint, 'log.txt'), title=title)
logger.set_names(
['Epoch', 'LR', 'Train Loss', 'Val Loss', 'Train Acc', 'Val Acc'])
print(' Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
# create data loader
train_dataset = datasets.__dict__[args.dataset](is_train=True,
**vars(args))
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
val_dataset = datasets.__dict__[args.dataset](is_train=False, **vars(args))
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# evaluation only
if args.evaluate:
print('\nEvaluation only')
loss, acc, predictions = validate(val_loader, model, criterion,
njoints, args.debug, args.flip)
save_pred(predictions, checkpoint=args.checkpoint)
return
# train and eval
lr = args.lr
for epoch in range(args.start_epoch, args.epochs):
lr = adjust_learning_rate(optimizer, epoch, lr, args.schedule,
args.gamma)
print('\nEpoch: %d | LR: %.8f' % (epoch + 1, lr))
# decay sigma
if args.sigma_decay > 0:
train_loader.dataset.sigma *= args.sigma_decay
val_loader.dataset.sigma *= args.sigma_decay
# train for one epoch
train_loss, train_acc = train(train_loader, model, criterion,
optimizer, args.debug, args.flip)
# evaluate on validation set
valid_loss, valid_acc, predictions = validate(val_loader, model,
criterion, njoints,
args.debug, args.flip)
# append logger file
logger.append(
[epoch + 1, lr, train_loss, valid_loss, train_acc, valid_acc])
# remember best acc and save checkpoint
is_best = valid_acc > best_acc
best_acc = max(valid_acc, best_acc)
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc': best_acc,
'optimizer': optimizer.state_dict(),
},
predictions,
is_best,
checkpoint=args.checkpoint,
snapshot=args.snapshot)
logger.close()
logger.plot(['Train Acc', 'Val Acc'])
savefig(os.path.join(args.checkpoint, 'log.eps'))
def main(args):
global best_acc
global idx1
global idx2
# idx is the index of joints used to compute accuracy for dataset2
if args.dataset1 == 'real_animal' or args.dataset1 == 'real_animal_sp' or args.dataset1 == 'synthetic_animal' or args.dataset1 == 'synthetic_animal_sp':
idx1 = range(1, 19)
else:
print("Unknown dataset: {}".format(args.dataset1))
assert False
if args.dataset2 == 'real_animal' or args.dataset2 == 'real_animal_sp':
if args.animal == 'horse':
idx2 = range(1, 19) # horse
elif args.animal == 'tiger':
idx2 = [
1, 2, 3, 4, 5, 6, 7, 8, 15, 16, 17, 18, 13, 14, 9, 10, 11, 12
] # tiger
else:
print("Unknown dataset: {}".format(args.dataset2))
assert False
# create checkpoint dir
if not isdir(args.checkpoint):
mkdir_p(args.checkpoint)
# create model
njoints = datasets.__dict__[args.dataset1].njoints
print(njoints)
print("==> creating model '{}', stacks={}, blocks={}".format(
args.arch, args.stacks, args.blocks))
model = models.__dict__[args.arch](num_stacks=args.stacks,
num_blocks=args.blocks,
num_classes=njoints,
resnet_layers=args.resnet_layers)
model = torch.nn.DataParallel(model).to(device)
# define loss function (criterion) and optimizer
criterion = losses.JointsMSELoss().to(device)
if args.solver == 'rms':
optimizer = torch.optim.RMSprop(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.solver == 'adam':
optimizer = torch.optim.Adam(
model.parameters(),
lr=args.lr,
)
else:
print('Unknown solver: {}'.format(args.solver))
assert False
# optionally resume from a checkpoint
title = args.dataset1 + ' ' + args.arch
if args.resume:
if isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
#model.load_state_dict(checkpoint)
args.start_epoch = checkpoint['epoch']
best_acc = checkpoint['best_acc']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
else:
raise Exception('please provide a checkpoint')
print(' Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
# create data loader
train_dataset = datasets.__dict__[args.dataset1](is_train=True,
**vars(args))
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.train_batch,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
val_dataset = datasets.__dict__[args.dataset2](is_train=False,
**vars(args))
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.test_batch,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# evaluation only
if args.evaluate:
print('\nEvaluation only')
_, acc, predictions = validate(val_loader, model, criterion, njoints,
args, args.flip, args.test_batch)
#save_pred(predictions, checkpoint=args.checkpoint)
return
def main(args):
global best_acc
global idx
global customMpiiObject
# idx is the index of joints used to compute accuracy
if args.dataset in ['mpii', 'lsp']:
idx = [1, 2, 3, 4, 5, 6, 11, 12, 15, 16]
elif args.dataset == 'coco':
idx = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17]
else:
print("Unknown dataset: {}".format(args.dataset))
assert False
# create checkpoint dir
if not isdir(args.checkpoint):
mkdir_p(args.checkpoint)
# create model
njoints = CustomMpii.njoints
print("==> creating model '{}', stacks={}, blocks={}".format(
args.arch, args.stacks, args.blocks))
model = models.__dict__[args.arch](num_stacks=args.stacks,
num_blocks=args.blocks,
num_classes=njoints,
resnet_layers=args.resnet_layers)
model = torch.nn.DataParallel(model).to(device)
# define loss function (criterion) and optimizer
criterion = losses.JointsMSELoss().to(device)
if args.solver == 'rms':
optimizer = torch.optim.RMSprop(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.solver == 'adam':
optimizer = torch.optim.Adam(
model.parameters(),
lr=args.lr,
)
else:
print('Unknown solver: {}'.format(args.solver))
assert False
# optionally resume from a checkpoint
title = args.dataset + ' ' + args.arch
if args.resume:
if isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume, map_location='cpu')
args.start_epoch = checkpoint['epoch']
best_acc = checkpoint['best_acc']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
logger = Logger(join(args.checkpoint, 'log.txt'),
title=title,
resume=True)
else:
print("=> no checkpoint found at '{}'".format(args.resume))
else:
logger = Logger(join(args.checkpoint, 'log.txt'), title=title)
logger.set_names(
['Epoch', 'LR', 'Train Loss', 'Val Loss', 'Train Acc', 'Val Acc'])
print(' Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
print("Loading validation dataset...")
customMpiiObject = CustomMpii(is_train=False, **vars(args))
loss, acc, predictions = validate(args.unit_path, model, criterion,
njoints, args.debug, args.flip)
print(predictions)
if args.debug:
paint.main(args.unit_path, predictions.squeeze().tolist())
return
from pose.utils.imutils import batch_with_heatmap
from pose.utils.transforms import fliplr, flip_back
import pose.models as models
import pose.losses as losses
import scripts.ssl_datasets as datasets
from options.train_options import TrainOptions
from CCSSL.scripts.timer import Timer
from CCSSL.scripts.consistency import prediction_check
opt = TrainOptions()
args = opt.initialize()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
idx = range(1, args.num_classes + 1)
njoints = len(idx)
criterion = losses.JointsMSELoss().to(device)
global_animal = args.animal
def main():
_t = {'iter time': Timer()}
# create directory
model_name = args.source + '_to_' + args.target
if not os.path.exists(args.checkpoint):
os.makedirs(args.checkpoint)
os.makedirs(os.path.join(args.checkpoint, 'logs'))
os.makedirs(os.path.join(args.checkpoint, 'ssl_labels'))
opt.print_options(args)
def main(args):
img_path = "/home/shantam/Documents/Programs/hourglasstensorlfow/images/cropped0.jpg"
img = load_image(img_path)
c = [img.shape[1] / 2, img.shape[2] / 2]
s = float(img.shape[1] / 200.0)
img = crop(img_path, img, c, s, [256, 256])
trans = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
#img = trans(img)
img = img.unsqueeze(0)
print(img.shape)
njoints = 24
model = models.__dict__[args.arch](num_stacks=args.stacks,
num_blocks=args.blocks,
num_classes=njoints,
resnet_layers=args.resnet_layers)
model = torch.nn.DataParallel(model).to(device)
criterion = losses.JointsMSELoss().to(device)
if args.solver == 'rms':
print("done")
optimizer = torch.optim.RMSprop(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.solver == 'adam':
optimizer = torch.optim.Adam(
model.parameters(),
lr=args.lr,
)
checkpoint = torch.load(
"/home/shantam/Documents/Programs/pytorch-pose/checkpoint/mpii/hg_fullset/checkpoint.pth.tar"
)
args.start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
model.eval()
img = img.to(device, non_blocking=True)
output = model(img)
#print (len(output))
score_map = output[-1].cpu() if type(output) == list else output.cpu()
preds, vals = final_preds(score_map, [c], [s], [64, 64])
image = cv2.imread(img_path)
coords = np.squeeze(preds)
for m in range(0, len(coords)):
val = vals[0][m].detach().numpy()
print(val)
if val > 0.25: #threshold for confidence score
x, y = coords[m][0].cpu().detach().numpy(), coords[m][1].cpu(
).detach().numpy()
print(x, y)
cv2.circle(image, (x, y), 1, (0, 0, 255), -1)
while True:
cv2.imshow("dec", image)
cv2.waitKey(10)