def setUp(self, mock):
# prepare mock.
mock.return_value = ['image1.png,108.0,50.0,1.0,148.0,180.0,0.0,148.0,180.0,1.0\n',
'image2.png,40.0,50.0,1.0,160.0,180.0,1.0,160.0,180.0,0.0\n']
# set up.
self.path = 'test_data'
self.dataset = PoseDataset(
self.path,
input_transform=transforms.Compose([
transforms.ToTensor()]),
output_transform=Scale(),
transform=Crop(data_augmentation=False))
def __init__(self, Nj, gpu, model_file, filename):
# validate arguments.
self.gpu = (gpu >= 0)
if self.gpu and not torch.cuda.is_available():
raise GPUNotFoundError('GPU is not found.')
# initialize model to estimate.
self.model = AlexNet(Nj)
self.model.load_state_dict(torch.load(model_file))
# prepare gpu.
if self.gpu:
self.model.cuda()
# load dataset to estimate.
self.dataset = PoseDataset(filename,
input_transform=transforms.Compose(
[transforms.ToTensor(),
RandomNoise()]),
output_transform=Scale(),
transform=Crop())
def setUp(self):
self.value = 227
self.transform = Scale(self.value)
def start(self):
""" Train pose net. """
# set random seed.
if self.seed is not None:
random.seed(self.seed)
torch.manual_seed(self.seed)
if self.gpu:
torch.cuda.manual_seed(self.seed)
# initialize model to train.
model = AlexNet(self.Nj)
if self.resume_model:
model.load_state_dict(torch.load(self.resume_model))
# prepare gpu.
if self.gpu:
model.cuda()
# load the datasets.
input_transforms = [transforms.ToTensor()]
if self.data_augmentation:
input_transforms.append(RandomNoise())
train = PoseDataset(
self.train,
input_transform=transforms.Compose(input_transforms),
output_transform=Scale(),
transform=Crop(data_augmentation=self.data_augmentation))
val = PoseDataset(self.val,
input_transform=transforms.Compose(
[transforms.ToTensor()]),
output_transform=Scale(),
transform=Crop(data_augmentation=False))
# training/validation iterators.
train_iter = torch.utils.data.DataLoader(train,
batch_size=self.batchsize,
shuffle=True)
val_iter = torch.utils.data.DataLoader(val,
batch_size=self.batchsize,
shuffle=False)
# set up an optimizer.
optimizer = self._get_optimizer(model)
if self.resume_opt:
optimizer.load_state_dict(torch.load(self.resume_opt))
# set intervals.
val_interval = 10
resume_interval = self.epoch / 10
log_interval = 10
# set logger and start epoch.
logger = TrainLogger(os.path.join(self.out, 'pytorch'))
start_epoch = 0
if self.resume:
resume = torch.load(self.resume)
start_epoch = resume['epoch']
logger.load_state_dict(resume['logger'])
# start training.
start_time = time.time()
for epoch in trange(start_epoch,
self.epoch,
initial=start_epoch,
total=self.epoch,
desc=' total'):
self._train(model, optimizer, train_iter, log_interval, logger,
start_time)
if (epoch + 1) % val_interval == 0:
self._test(model, val_iter, logger, start_time)
if (epoch + 1) % resume_interval == 0:
self._checkpoint(epoch, model, optimizer, logger)
def __init__(self,
Nj,
NN,
gpu,
model_file,
filename,
Dataset3D,
isEval=True):
# validate arguments.
self.gpu = (gpu >= 0)
self.NN = NN
if self.gpu and not torch.cuda.is_available():
raise GPUNotFoundError('GPU is not found.')
# initialize model to estimate.
if self.NN == "MobileNet_":
self.model = MobileNet_()
elif self.NN == "MobileNet__":
self.model = MobileNet__()
elif self.NN == "MobileNet_3":
self.model = MobileNet_3()
elif self.NN == "MobileNet_4":
self.model = MobileNet_4()
elif self.NN == "MobileNet___":
self.model = MobileNet___()
elif self.NN == "MobileNet":
self.model = MobileNet()
elif self.NN == "MnasNet":
self.model = MnasNet()
elif self.NN == "MnasNet_":
self.model = MnasNet_()
elif self.NN == "MnasNet56_":
self.model = MnasNet56_()
elif self.NN == "AlexNet":
self.model = AlexNet(Nj)
elif self.NN == "MobileNet3D":
self.model = MobileNet3D()
elif self.NN == "MnasNet3D":
self.model = MnasNet3D()
elif self.NN == "MobileNet3D2":
self.model = MobileNet3D2()
else:
self.model = Resnet()
self.model.load_state_dict(torch.load(model_file))
if isEval == True:
self.model.eval()
# prepare gpu.
if self.gpu:
self.model.cuda()
# load dataset to estimate.
if Dataset3D:
'''
self.dataset = PoseDataset3D(
filename,
input_transform=transforms.ToTensor())
'''
self.dataset = PoseDataset3D(filename,
input_transform=transforms.Compose([
transforms.ToTensor(),
RandomNoise()
]))
else:
self.dataset = PoseDataset(filename,
input_transform=transforms.Compose([
transforms.ToTensor(),
RandomNoise()
]),
output_transform=Scale(),
transform=Crop(data_augmentation=False))
def main():
""" Main function. """
# arg definition
parser = argparse.ArgumentParser(
description='Visualize image and its pose in dataset.')
parser.add_argument('path',
type=str,
help='Path to dataset (image-pose list file).')
parser.add_argument('--out',
default='result/dataset',
help='Output directory.')
parser.add_argument('--use-visibility',
'-v',
action='store_true',
help='Use visibility to plot pose.')
args = parser.parse_args()
output_dir = os.path.join(args.out, os.path.basename(args.path))
# create directory.
try:
os.makedirs(output_dir)
except OSError:
pass
# get dataset.
dataset = PoseDataset(args.path,
input_transform=transforms.Compose(
[transforms.ToTensor()]),
output_transform=Scale(),
transform=Crop(data_augmentation=False))
for index, (image, pose, visibility,
image_types) in enumerate(tqdm(dataset, ascii=True)):
# get data.
_, size, _ = image.shape
'''
if args.use_visibility:
pose = pose[visibility.ravel().astype(bool)]
else:
pose = pose.abs()
pose *= size
pose_x, pose_y = zip(*pose)
'''
scale = float(size) / float(14)
if args.use_visibility:
pose = pose[visibility.ravel().astype(bool)]
else:
pose = pose.abs()
#pose = (pose*scale + 0.5).int() * 14
pose = (pose * 14).int() * 16
pose_x, pose_y = zip(*pose)
# plot image and pose.
fig = plt.figure(figsize=(2.24, 2.24))
img = image.numpy().transpose(1, 2, 0)
plt.imshow(img, vmin=0., vmax=1.)
for i in range(14):
plt.scatter(pose_x[i],
pose_y[i],
color=cm.hsv(i / 14.0),
s=17 - (i + 2))
plt.axis("off")
plt.savefig(
os.path.join(output_dir, '{0}{1}.png'.format(image_types, index)))
plt.close(fig)