Add code of GaitBase (#115)

* add resnet9 backbone and regular da ops

* add gait3d config

* fix invalid path CASIA-B* in windows

* add gaitbase config for all datasets

* rm unused OpenGait transform

opengait/data/transform.py

@@ -1,6 +1,9 @@
-from data import transform as base_transform
 import numpy as np
-
+import random
+import torchvision.transforms as T
+import cv2
+import math
+from data import transform as base_transform
 from utils import is_list, is_dict, get_valid_args
 
 
@@ -49,6 +52,139 @@ class BaseRgbTransform():
         return (x - self.mean) / self.std
 
 
+# **************** Data Agumentation ****************
+
+
+class RandomHorizontalFlip(object):
+    def __init__(self, prob=0.5):
+        self.prob = prob
+
+    def __call__(self, seq):
+        if random.uniform(0, 1) >= self.prob:
+            return seq
+        else:
+            return seq[:, :, ::-1]
+
+
+class RandomErasing(object):
+    def __init__(self, prob=0.5, sl=0.05, sh=0.2, r1=0.3, per_frame=False):
+        self.prob = prob
+        self.sl = sl
+        self.sh = sh
+        self.r1 = r1
+        self.per_frame = per_frame
+
+    def __call__(self, seq):
+        if not self.per_frame:
+            if random.uniform(0, 1) >= self.prob:
+                return seq
+            else:
+                for _ in range(100):
+                    seq_size = seq.shape
+                    area = seq_size[1] * seq_size[2]
+
+                    target_area = random.uniform(self.sl, self.sh) * area
+                    aspect_ratio = random.uniform(self.r1, 1 / self.r1)
+
+                    h = int(round(math.sqrt(target_area * aspect_ratio)))
+                    w = int(round(math.sqrt(target_area / aspect_ratio)))
+
+                    if w < seq_size[2] and h < seq_size[1]:
+                        x1 = random.randint(0, seq_size[1] - h)
+                        y1 = random.randint(0, seq_size[2] - w)
+                        seq[:, x1:x1+h, y1:y1+w] = 0.
+                        return seq
+            return seq
+        else:
+            self.per_frame = False
+            frame_num = seq.shape[0]
+            ret = [self.__call__(seq[k][np.newaxis, ...])
+                   for k in range(frame_num)]
+            self.per_frame = True
+            return np.concatenate(ret, 0)
+
+
+class RandomRotate(object):
+    def __init__(self, prob=0.5, degree=10):
+        self.prob = prob
+        self.degree = degree
+
+    def __call__(self, seq):
+        if random.uniform(0, 1) >= self.prob:
+            return seq
+        else:
+            _, dh, dw = seq.shape
+            # rotation
+            degree = random.uniform(-self.degree, self.degree)
+            M1 = cv2.getRotationMatrix2D((dh // 2, dw // 2), degree, 1)
+            # affine
+            seq = [cv2.warpAffine(_[0, ...], M1, (dw, dh))
+                   for _ in np.split(seq, seq.shape[0], axis=0)]
+            seq = np.concatenate([np.array(_)[np.newaxis, ...]
+                                 for _ in seq], 0)
+            return seq
+
+
+class RandomPerspective(object):
+    def __init__(self, prob=0.5):
+        self.prob = prob
+
+    def __call__(self, seq):
+        if random.uniform(0, 1) >= self.prob:
+            return seq
+        else:
+            _, h, w = seq.shape
+            cutting = int(w // 44) * 10
+            x_left = list(range(0, cutting))
+            x_right = list(range(w - cutting, w))
+            TL = (random.choice(x_left), 0)
+            TR = (random.choice(x_right), 0)
+            BL = (random.choice(x_left), h)
+            BR = (random.choice(x_right), h)
+            srcPoints = np.float32([TL, TR, BR, BL])
+            canvasPoints = np.float32([[0, 0], [w, 0], [w, h], [0, h]])
+            perspectiveMatrix = cv2.getPerspectiveTransform(
+                np.array(srcPoints), np.array(canvasPoints))
+            seq = [cv2.warpPerspective(_[0, ...], perspectiveMatrix, (w, h))
+                   for _ in np.split(seq, seq.shape[0], axis=0)]
+            seq = np.concatenate([np.array(_)[np.newaxis, ...]
+                                 for _ in seq], 0)
+            return seq
+
+
+class RandomAffine(object):
+    def __init__(self, prob=0.5, degree=10):
+        self.prob = prob
+        self.degree = degree
+
+    def __call__(self, seq):
+        if random.uniform(0, 1) >= self.prob:
+            return seq
+        else:
+            _, dh, dw = seq.shape
+            # rotation
+            max_shift = int(dh // 64 * 10)
+            shift_range = list(range(0, max_shift))
+            pts1 = np.float32([[random.choice(shift_range), random.choice(shift_range)], [
+                              dh-random.choice(shift_range), random.choice(shift_range)], [random.choice(shift_range), dw-random.choice(shift_range)]])
+            pts2 = np.float32([[random.choice(shift_range), random.choice(shift_range)], [
+                              dh-random.choice(shift_range), random.choice(shift_range)], [random.choice(shift_range), dw-random.choice(shift_range)]])
+            M1 = cv2.getAffineTransform(pts1, pts2)
+            # affine
+            seq = [cv2.warpAffine(_[0, ...], M1, (dw, dh))
+                   for _ in np.split(seq, seq.shape[0], axis=0)]
+            seq = np.concatenate([np.array(_)[np.newaxis, ...]
+                                 for _ in seq], 0)
+            return seq
+        
+# ******************************************
+
+def Compose(trf_cfg):
+    assert is_list(trf_cfg)
+    transform = T.Compose([get_transform(cfg) for cfg in trf_cfg])
+    return transform
+
+
 def get_transform(trf_cfg=None):
     if is_dict(trf_cfg):
         transform = getattr(base_transform, trf_cfg['type'])

opengait/evaluation/evaluator.py

@@ -231,7 +231,7 @@ def evaluate_segmentation(data, dataset):
     return {"scalar/test_accuracy/mIOU": miou}
 
 
-def evaluate_Gait3D(data, conf, metric='euc'):
+def evaluate_Gait3D(data, dataset, metric='euc'):
     msg_mgr = get_msg_mgr()
 
     features, labels, cams, time_seqs = data['embeddings'], data['labels'], data['types'], data['views']

opengait/modeling/backbones/resnet.py

@@ -0,0 +1,58 @@
+from torch.nn import functional as F
+import torch.nn as nn
+from torchvision.models.resnet import BasicBlock, Bottleneck, ResNet
+from ..modules import BasicConv2d
+
+
+block_map = {'BasicBlock': BasicBlock,
+             'Bottleneck': Bottleneck}
+
+
+class ResNet9(ResNet):
+    def __init__(self, block, channels=[32, 64, 128, 256], in_channel=1, layers=[1, 2, 2, 1], strides=[1, 2, 2, 1], maxpool=True):
+        if block in block_map.keys():
+            block = block_map[block]
+        else:
+            raise ValueError(
+                "Error type for -block-Cfg-, supported: 'BasicBlock' or 'Bottleneck'.")
+        self.maxpool_flag = maxpool
+        super(ResNet9, self).__init__(block, layers)
+
+        # Not used #
+        self.fc = None
+        ############
+        self.inplanes = channels[0]
+        self.bn1 = nn.BatchNorm2d(self.inplanes)
+
+        self.conv1 = BasicConv2d(in_channel, self.inplanes, 3, 1, 1)
+
+        self.layer1 = self._make_layer(
+            block, channels[0], layers[0], stride=strides[0], dilate=False)
+
+        self.layer2 = self._make_layer(
+            block, channels[1], layers[1], stride=strides[1], dilate=False)
+        self.layer3 = self._make_layer(
+            block, channels[2], layers[2], stride=strides[2], dilate=False)
+        self.layer4 = self._make_layer(
+            block, channels[3], layers[3], stride=strides[3], dilate=False)
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
+        if blocks >= 1:
+            layer = super()._make_layer(block, planes, blocks, stride=stride, dilate=dilate)
+        else:
+            def layer(x): return x
+        return layer
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        if self.maxpool_flag:
+            x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        return x
+

opengait/modeling/base_model.py

@@ -427,7 +427,8 @@ class BaseModel(MetaModel, nn.Module):
                     model.train()
                     if model.cfgs['trainer_cfg']['fix_BN']:
                         model.fix_BN()
-                    model.msg_mgr.write_to_tensorboard(result_dict)
+                    if result_dict:
+                        model.msg_mgr.write_to_tensorboard(result_dict)
                     model.msg_mgr.reset_time()
             if model.iteration >= model.engine_cfg['total_iter']:
                 break