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)
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 test_load_dataset(self, mock): # prepare mock. mock.return_value = ['image1.png,10.0,20.0,0.0,30.0,40.0,1.0\n', 'image2.png,50.0,60.0,1.0,70.0,80.0,0.0\n'] # test. dataset = PoseDataset('test_data') eq_(dataset.images, ['image1.png', 'image2.png']) correct = [torch.Tensor([[10, 20], [30, 40]]), torch.Tensor([[50, 60], [70, 80]])] for p, c in zip(dataset.poses, correct): ok_((p == c).all()) correct = [torch.Tensor([[0, 0], [1, 1]]), torch.Tensor([[1, 1], [0, 0]])] for v, c in zip(dataset.visibilities, correct): ok_((v == c).all())
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())
class TestPoseDataset(unittest.TestCase): @patch('modules.dataset_indexing.pytorch.pose_dataset.open') 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) @patch('PIL.Image.open', return_value=Image.new('RGB', (320, 240))) def test_read_image(self, mock): image = self.dataset._read_image('dummy.png') eq_(type(image), Image.Image) eq_(image.size, (320, 240)) eq_(image.mode, 'RGB') def test_len(self): eq_(len(self.dataset), 2) @patch('PIL.Image.open') def test_getitem(self, mock): # prepare mock. shape = (256, 256) mock.side_effect = [Image.new('RGB', shape), Image.new('RGB', shape)] # test. image, pose, visibility = self.dataset[0] eq_(type(image), Image.Image) eq_(image.size, (256, 256)) eq_(image.mode, 'RGB') eq_(type(pose), torch.FloatTensor) ok_((pose == torch.Tensor([[108, 50], [148, 180], [148, 180]])).all()) eq_(type(visibility), torch.FloatTensor) ok_((visibility == torch.Tensor([[1, 1], [0, 0], [1, 1]])).all()) image, pose, visibility = self.dataset[1] eq_(type(image), Image.Image) eq_(image.size, (256, 256)) eq_(image.mode, 'RGB') eq_(type(pose), torch.FloatTensor) ok_((pose == torch.Tensor([[40, 50], [160, 180], [160, 180]])).all()) eq_(type(visibility), torch.FloatTensor) ok_((visibility == torch.Tensor([[1, 1], [1, 1], [0, 0]])).all())
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)