first commit

environment.yml

@@ -0,0 +1,51 @@
+name: gridconv
+channels:
+  - pytorch
+  - conda-forge
+  - defaults
+dependencies:
+  - blas=1.0=mkl
+  - configparser=4.0.2=py27_0
+  - cudatoolkit=10.1.243=h036e899_8
+  - hdf5=1.10.2=hba1933b_1
+  - intel-openmp=2022.0.1=h06a4308_3633
+  - ipykernel=4.10.0=py27_0
+  - ipython_genutils=0.2.0=pyhd3eb1b0_1
+  - ipywidgets=7.6.0=pyhd3eb1b0_1
+  - libopencv=3.4.2=hb342d67_1
+  - mkl=2020.2=256
+  - mkl-service=2.3.0=py27he904b0f_0
+  - mkl_fft=1.0.15=py27ha843d7b_0
+  - mkl_random=1.1.0=py27hd6b4f25_0
+  - ninja=1.10.2=h5e70eb0_2
+  - numpy=1.16.6=py27hbc911f0_0
+  - numpy-base=1.16.6=py27hde5b4d6_0
+  - opencv=3.4.2=py27h6fd60c2_1
+  - pillow=6.2.1=py27h34e0f95_0
+  - pip=19.3.1=py27_0
+  - py-opencv=3.4.2=py27hb342d67_1
+  - pycparser=2.20=py_2
+  - pyparsing=2.4.7=pyhd3eb1b0_0
+  - python=2.7.18=ha1903f6_2
+  - python-dateutil=2.8.2=pyhd3eb1b0_0
+  - pytorch=1.4.0=py2.7_cuda10.1.243_cudnn7.6.3_0
+  - six=1.16.0=pyhd3eb1b0_1
+  - torchaudio=0.4.0=py27
+  - torchvision=0.5.0=py27_cu101
+  - wheel=0.37.1=pyhd3eb1b0_0
+  - widgetsnbextension=3.5.1=py27_0
+  - zlib=1.2.11=h7f8727e_4
+  - pip:
+    - ipdb==0.13.9
+    - ipython==5.10.0
+    - kiwisolver==1.1.0
+    - matplotlib==2.2.5
+    - pickleshare==0.7.5
+    - progress==1.6
+    - protobuf==3.17.3
+    - pyyaml==5.4.1
+    - scikit-learn==0.20.4
+    - scipy==1.2.3
+    - sklearn==0.0
+    - subprocess32==3.5.4
+    - tqdm==4.63.0

src/base_modules.py

@@ -0,0 +1,68 @@
+import torch
+import torch.nn as nn
+import numpy as np
+import os
+
+from network.gridconv import GridLiftingNetwork
+from network.dgridconv import DynamicGridLiftingNetwork
+from network.dgridconv_autogrids import AutoDynamicGridLiftingNetwork
+from dataset.human36m import Human36M
+
+def get_dataloader(opt, is_train=False, shuffle=False):
+    if not is_train and opt.input != 'gt':
+        exclude_drift_data = True
+    else:
+        exclude_drift_data = False
+    actual_data_dir = os.path.join(opt.data_rootdir, opt.input)
+
+    dataset = Human36M(data_path=actual_data_dir, is_train=is_train, exclude_drift_data=exclude_drift_data, prepare_grid=opt.prepare_grid)
+
+    dataloader = torch.utils.data.DataLoader(
+        dataset=dataset,
+        batch_size=opt.batch if is_train else opt.test_batch,
+        shuffle=shuffle,
+        num_workers=0,
+        pin_memory=True
+    )
+    return dataloader
+
+def get_lifting_model(opt):
+    if opt.lifting_model == 'gridconv':
+        model = GridLiftingNetwork(hidden_size=opt.hidsize,
+                                   num_block=opt.num_block)
+    elif opt.lifting_model == 'dgridconv':
+        model = DynamicGridLiftingNetwork(hidden_size=opt.hidsize,
+                                          num_block=opt.num_block,
+                                          grid_shape=opt.grid_shape,
+                                          padding_mode=opt.padding_mode)
+    elif opt.lifting_model == 'dgridconv_autogrids':
+        model = AutoDynamicGridLiftingNetwork(hidden_size=opt.hidsize,
+                                          num_block=opt.num_block,
+                                          grid_shape=opt.grid_shape,
+                                          padding_mode=opt.padding_mode)
+    else:
+        raise Exception('Unexpected argument, %s' % opt.lifting_model)
+    model = model.cuda()
+    if opt.load:
+        ckpt = torch.load(opt.load)
+        model.load_state_dict(ckpt['state_dict'])
+    return model
+
+def get_optimizer(model, opt):
+    optimizer = torch.optim.Adam(model.parameters(), lr=opt.lr)
+    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=opt.lr_gamma)
+
+    return optimizer, scheduler
+
+
+def get_loss(opt):
+    if opt.loss == 'l2':
+        criterion = nn.MSELoss(reduction='mean').cuda()
+    elif opt.loss == 'sqrtl2':
+        criterion = lambda output, target: torch.mean(torch.norm(output - target, dim=-1))
+    elif opt.loss == 'l1':
+        criterion = nn.L1Loss(reduction='mean').cuda()
+    else:
+        raise Exception('Unknown loss type %s' % opt.loss)
+
+    return criterion

src/dataset/human36m.py

@@ -0,0 +1,80 @@
+import os
+import torch
+import numpy as np
+from torch.utils.data import Dataset
+from tool import util
+import sys
+from tqdm import tqdm
+
+
+S9_drift_fname_list = ['Waiting 1.60457274', 'Greeting.60457274', 'Greeting.58860488', 'SittingDown 1.58860488',
+                        'Waiting 1.54138969', 'SittingDown 1.54138969', 'Waiting 1.55011271', 'Greeting.54138969',
+                        'Greeting.55011271', 'SittingDown 1.60457274', 'SittingDown 1.55011271', 'Waiting 1.58860488']
+
+
+class Human36M(Dataset):
+    def __init__(self, data_path, is_train, exclude_drift_data, prepare_grid):
+        self.data_path = data_path
+        self.exclude_drift_data = exclude_drift_data
+        self.num_jts = 17
+        self.inp, self.out = [], []
+        self.meta = {'info':[]}
+        self.confidence_2d = []
+        self.subject_list = ['S1','S5','S6','S7','S8']  if is_train else ['S9', 'S11']
+        self.prepare_grid = prepare_grid
+
+        data_2d = {}
+        data_3d = {}
+        self.phase = 'train' if is_train else 'test'
+
+        for data_prefix in [self.phase]:
+            data_2d_file = '%s_custom_2d_unnorm.pth.tar' % data_prefix
+            data_3d_file = '%s_custom_3d_unnorm.pth.tar' % data_prefix
+            cur_data_2d = torch.load(os.path.join(data_path, data_2d_file))
+            cur_data_3d = torch.load(os.path.join(data_path, data_3d_file))
+            data_2d.update(cur_data_2d)
+            data_3d.update(cur_data_3d)
+
+        ordered_key = sorted(data_2d.keys())
+        ordered_key = list(filter(lambda x: x[0] in self.subject_list, ordered_key))
+        sample_step = 1
+        for key in tqdm(ordered_key):
+            sub, act, fname = key
+            fullact = fname.split('.')[0]
+            num_f = data_2d[key].shape[0]
+            if (sub == 'S11') and (fullact == 'Directions'):
+                continue
+            if self.exclude_drift_data and sub == 'S9' and fname in S9_drift_fname_list:
+                continue
+            for i in range(0, num_f, sample_step):
+                p2d_ori = data_2d[key][i].reshape(self.num_jts, 2)
+                p3d_ori = data_3d[key]['joint_3d'][i].reshape(self.num_jts, 3)
+
+                p2d = (p2d_ori - 500) / 500.
+                p3d = p3d_ori / 1000.
+                self.inp.append(p2d)
+                self.out.append(p3d)
+                self.meta['info'].append({'subject':sub, 'action':fullact, 'camid':fname.split('.')[-1], 'frid':i})
+
+
+
+    def __getitem__(self, index):
+        inputs = self.inp[index].copy()
+        outputs = self.out[index].copy()
+
+        if self.prepare_grid:
+            inputs = util.semantic_grid_trans(np.expand_dims(inputs, axis=0)).squeeze(0)
+            outputs = util.semantic_grid_trans(np.expand_dims(outputs, axis=0)).squeeze(0)
+
+        inputs = torch.Tensor(inputs).float()
+        outputs = torch.Tensor(outputs).float()
+
+        meta = self.meta['info'][index]
+        for key in self.meta:
+            if key != 'info':
+                meta[key] = self.meta[key]
+
+        return inputs, outputs, meta
+
+    def __len__(self):
+        return len(self.inp)

src/lifting.py

@@ -0,0 +1,143 @@
+import torch
+import time
+from progress.bar import Bar
+import numpy as np
+
+from tool import util
+
+def train(epoch, train_loader, lifting_model, criterion, optimizer, opt):
+    losses = util.AverageMeter()
+
+    start = time.time()
+    batch_time = 0
+    train_loader_len = len(train_loader)
+    bar = Bar('>>>', fill='>', max=train_loader_len)
+
+    lifting_model.train()
+
+    for i, (inputs, targets, _) in enumerate(train_loader):
+        batch_size = targets.shape[0]
+
+        if hasattr(lifting_model, "net_update_temperature") and epoch < opt.temp_epoch:
+            temperature = util.get_temperature(0, epoch, opt.temp_epoch, i, train_loader_len,
+                                          method='linear', max_temp=opt.max_temp, increase=False)
+            lifting_model.net_update_temperature(temperature)
+
+        inputs_gpu = inputs.cuda()
+        targets_gpu = targets.cuda()
+
+        outputs_gpu = lifting_model(inputs_gpu)
+        optimizer.zero_grad()
+        loss = criterion(outputs_gpu, targets_gpu)
+        losses.update(loss.item(), batch_size)
+        loss.backward()
+
+        optimizer.step()
+
+        # update summary
+        if (i + 1) % 100 == 0:
+            batch_time = time.time() - start
+            start = time.time()
+
+        bar.suffix = '({batch}/{size}) | batch: {batchtime:.1}ms | Total: {ttl} | ETA: {eta:} | loss: {loss:.6f}' \
+            .format(batch=i + 1,
+                    size=len(train_loader),
+                    batchtime=batch_time * 10.0,
+                    ttl=bar.elapsed_td,
+                    eta=bar.eta_td,
+                    loss=losses.avg)
+        bar.next()
+    bar.finish()
+    return losses.avg
+
+def evaluate(test_loader, lifting_model, criterion, opt):
+    loss_test, outputs_array, targets_array, corresponding_info = inference(test_loader, lifting_model, criterion, opt)
+    num_sample = len(outputs_array)
+
+    outputs_array_by_action = util.rearrange_by_key(outputs_array, corresponding_info)
+    targets_array_by_action = util.rearrange_by_key(targets_array, corresponding_info)
+
+    err_ttl, err_act, err_dim = evaluate_actionwise(outputs_array_by_action, targets_array_by_action, opt.procrustes)
+
+    print(">>> error mean of %d samples: %.3f <<<" % (num_sample, err_ttl))
+    print(">>> error by dim: x: %.3f,  y:%.3f, z:%.3f <<<" % (tuple(err_dim)))
+    return loss_test, err_ttl, err_dim
+
+def inference(test_loader, lifting_model, criterion, opt):
+    print('Inferring...')
+    losses = util.AverageMeter()
+    lifting_model.eval()
+    outputs_array = []
+    targets_array = []
+    corresponding_info = []
+
+    start = time.time()
+    batch_time = 0
+    bar = Bar('>>>', fill='>', max=len(test_loader))
+
+    with torch.no_grad():
+        for i, (inputs, targets, meta) in enumerate(test_loader):
+            batch_size = targets.shape[0]
+            inputs_gpu = inputs.cuda()
+            info = meta
+            info['fullaction'] = meta['action']
+            info['action'] = list(map(lambda x: x.split(' ')[0], meta['action']))
+            outputs_gpu = lifting_model(inputs_gpu)
+            targets_gpu = targets.cuda()
+
+            loss = criterion(outputs_gpu, targets_gpu)
+            losses.update(loss.item(), batch_size)
+
+            if opt.prepare_grid:
+                outputs_pose = util.inverse_semantic_grid_trans(outputs_gpu.cpu().data.numpy())
+                targets_pose = util.inverse_semantic_grid_trans(targets.data.numpy())
+            else:
+                outputs_pose = outputs_gpu.cpu().data.numpy()
+                targets_pose = targets.data.numpy()
+
+            outputs_array.append(outputs_pose)
+            targets_array.append(targets_pose)
+            info_list = util.dict2list(info)
+            corresponding_info += info_list
+
+            bar.suffix = '({batch}/{size}) | batch: {batchtime:.1}ms | Total: {ttl} | ETA: {eta:} | loss: {loss:.6f}' \
+                .format(batch=i + 1,
+                        size=len(test_loader),
+                        batchtime=batch_time * 10.0,
+                        ttl=bar.elapsed_td,
+                        eta=bar.eta_td,
+                        loss=losses.avg)
+            bar.next()
+    bar.finish()
+
+    outputs_array = np.vstack(outputs_array)
+    targets_array = np.vstack(targets_array)
+    return losses.avg, outputs_array, targets_array, corresponding_info
+
+def evaluate_actionwise(outputs_array, targets_array, procrustes):
+    err_ttl = util.AverageMeter()
+    err_act = {}
+    err_dim = [util.AverageMeter(), util.AverageMeter(), util.AverageMeter()]
+
+    for act in sorted(outputs_array.keys()):
+        num_sample = outputs_array[act].shape[0]
+        predict = outputs_array[act] * 1000.
+        gt = targets_array[act] * 1000.
+
+        if procrustes:
+            pred_procrustes = []
+            for i in range(num_sample):
+                _, Z, T, b, c = util.get_procrustes_transformation(gt[i], predict[i], True)
+                pred_procrustes.append((b * predict[i].dot(T)) + c)
+            predict = np.array(pred_procrustes)
+
+        err_act[act] = (((predict - gt) ** 2).sum(-1)**0.5).mean()
+        err_ttl.update(err_act[act], 1)
+        for dim_i in range(len(err_dim)):
+            err = (np.abs(predict[:, :, dim_i] - gt[:, :, dim_i])).mean()
+            err_dim[dim_i].update(err, 1)
+
+    for dim_i in range(len(err_dim)):
+        err_dim[dim_i] = err_dim[dim_i].avg
+
+    return err_ttl.avg, err_act, err_dim