ADD file via upload

Gaitset

@@ -0,0 +1,192 @@
+
+import os
+from scipy import misc as scisc
+import cv2
+import numpy as np
+from warnings import warn
+from time import sleep
+import argparse
+
+from multiprocessing import Pool
+from multiprocessing import TimeoutError as MP_TimeoutError
+
+START = "START"
+FINISH = "FINISH"
+WARNING = "WARNING"
+FAIL = "FAIL"
+
+
+def boolean_string(s):
+    if s.upper() not in {'FALSE', 'TRUE'}:
+        raise ValueError('Not a valid boolean string')
+    return s.upper() == 'TRUE'
+
+
+parser = argparse.ArgumentParser(description='Test')
+parser.add_argument('--input_path', default='', type=str,
+                    help='Root path of raw dataset.')
+parser.add_argument('--output_path', default='', type=str,
+                    help='Root path for output.')
+parser.add_argument('--log_file', default='./pretreatment.log', type=str,
+                    help='Log file path. Default: ./pretreatment.log')
+parser.add_argument('--log', default=False, type=boolean_string,
+                    help='If set as True, all logs will be saved. '
+                         'Otherwise, only warnings and errors will be saved.'
+                         'Default: False')
+parser.add_argument('--worker_num', default=1, type=int,
+                    help='How many subprocesses to use for data pretreatment. '
+                         'Default: 1')
+opt = parser.parse_args()
+
+INPUT_PATH = opt.input_path
+OUTPUT_PATH = opt.output_path
+IF_LOG = opt.log
+LOG_PATH = opt.log_file
+WORKERS = opt.worker_num
+
+T_H = 64
+T_W = 64
+
+
+def log2str(pid, comment, logs):
+    str_log = ''
+    if type(logs) is str:
+        logs = [logs]
+    for log in logs:
+        str_log += "# JOB %d : --%s-- %s\n" % (
+            pid, comment, log)
+    return str_log
+
+
+def log_print(pid, comment, logs):
+    str_log = log2str(pid, comment, logs)
+    if comment in [WARNING, FAIL]:
+        with open(LOG_PATH, 'a') as log_f:
+            log_f.write(str_log)
+    if comment in [START, FINISH]:
+        if pid % 500 != 0:
+            return
+    print(str_log, end='')
+
+
+def cut_img(img, seq_info, frame_name, pid):
+    # A silhouette contains too little white pixels
+    # might be not valid for identification.
+    if img.sum() <= 10000:
+        message = 'seq:%s, frame:%s, no data, %d.' % (
+            '-'.join(seq_info), frame_name, img.sum())
+        warn(message)
+        log_print(pid, WARNING, message)
+        return None
+    # Get the top and bottom point
+    y = img.sum(axis=1)
+    y_top = (y != 0).argmax(axis=0)
+    y_btm = (y != 0).cumsum(axis=0).argmax(axis=0)
+    img = img[y_top:y_btm + 1, :]
+    # As the height of a person is larger than the width,
+    # use the height to calculate resize ratio.
+    _r = img.shape[1] / img.shape[0]
+    _t_w = int(T_H * _r)
+    img = cv2.resize(img, (_t_w, T_H), interpolation=cv2.INTER_CUBIC)
+    # Get the median of x axis and regard it as the x center of the person.
+    sum_point = img.sum()
+    sum_column = img.sum(axis=0).cumsum()
+    x_center = -1
+    for i in range(sum_column.size):
+        if sum_column[i] > sum_point / 2:
+            x_center = i
+            break
+    if x_center < 0:
+        message = 'seq:%s, frame:%s, no center.' % (
+            '-'.join(seq_info), frame_name)
+        warn(message)
+        log_print(pid, WARNING, message)
+        return None
+    h_T_W = int(T_W / 2)
+    left = x_center - h_T_W
+    right = x_center + h_T_W
+    if left <= 0 or right >= img.shape[1]:
+        left += h_T_W
+        right += h_T_W
+        _ = np.zeros((img.shape[0], h_T_W))
+        img = np.concatenate([_, img, _], axis=1)
+    img = img[:, left:right]
+    return img.astype('uint8')
+
+
+def cut_pickle(seq_info, pid):
+    seq_name = '-'.join(seq_info)
+    log_print(pid, START, seq_name)
+    seq_path = os.path.join(INPUT_PATH, *seq_info)
+    out_dir = os.path.join(OUTPUT_PATH, *seq_info)
+    frame_list = os.listdir(seq_path)
+    frame_list.sort()
+    count_frame = 0
+    for _frame_name in frame_list:
+        frame_path = os.path.join(seq_path, _frame_name)
+        img = cv2.imread(frame_path)[:, :, 0]
+        img = cut_img(img, seq_info, _frame_name, pid)
+        if img is not None:
+            # Save the cut img
+            save_path = os.path.join(out_dir, _frame_name)
+            scisc.imsave(save_path, img)
+            count_frame += 1
+    # Warn if the sequence contains less than 5 frames
+    if count_frame < 5:
+        message = 'seq:%s, less than 5 valid data.' % (
+            '-'.join(seq_info))
+        warn(message)
+        log_print(pid, WARNING, message)
+
+    log_print(pid, FINISH,
+              'Contain %d valid frames. Saved to %s.'
+              % (count_frame, out_dir))
+
+
+pool = Pool(WORKERS)
+results = list()
+pid = 0
+
+print('Pretreatment Start.\n'
+      'Input path: %s\n'
+      'Output path: %s\n'
+      'Log file: %s\n'
+      'Worker num: %d' % (
+          INPUT_PATH, OUTPUT_PATH, LOG_PATH, WORKERS))
+
+id_list = os.listdir(INPUT_PATH)
+id_list.sort()
+# Walk the input path
+for _id in id_list:
+    seq_type = os.listdir(os.path.join(INPUT_PATH, _id))
+    seq_type.sort()
+    for _seq_type in seq_type:
+        view = os.listdir(os.path.join(INPUT_PATH, _id, _seq_type))
+        view.sort()
+        for _view in view:
+            seq_info = [_id, _seq_type, _view]
+            out_dir = os.path.join(OUTPUT_PATH, *seq_info)
+            os.makedirs(out_dir)
+            results.append(
+                pool.apply_async(
+                    cut_pickle,
+                    args=(seq_info, pid)))
+            sleep(0.02)
+            pid += 1
+
+pool.close()
+unfinish = 1
+while unfinish > 0:
+    unfinish = 0
+    for i, res in enumerate(results):
+        try:
+            res.get(timeout=0.1)
+        except Exception as e:
+            if type(e) == MP_TimeoutError:
+                unfinish += 1
+                continue
+            else:
+                print('\n\n\nERROR OCCUR: PID ##%d##, ERRORTYPE: %s\n\n\n',
+                      i, type(e))
+                raise e
+pool.join()

READ ME.md

@@ -0,0 +1,157 @@
+# GaitSet
+
+[![LICENSE](https://img.shields.io/badge/license-NPL%20(The%20996%20Prohibited%20License)-blue.svg)](https://github.com/996icu/996.ICU/blob/master/LICENSE)
+[![996.icu](https://img.shields.io/badge/link-996.icu-red.svg)](https://996.icu)
+
+GaitSet is a **flexible**, **effective** and **fast** network for cross-view gait recognition. The [paper](https://ieeexplore.ieee.org/document/9351667) has been published on IEEE TPAMI.
+
+#### Flexible 
+The input of GaitSet is a set of silhouettes. 
+
+- There are **NOT ANY constrains** on an input,
+which means it can contain **any number** of **non-consecutive** silhouettes filmed under **different viewpoints**
+with **different walking conditions**.
+
+- As the input is a set, the **permutation** of the elements in the input
+will **NOT change** the output at all.
+
+#### Effective
+It achieves **Rank@1=95.0%** on [CASIA-B](http://www.cbsr.ia.ac.cn/english/Gait%20Databases.asp) 
+and  **Rank@1=87.1%** on [OU-MVLP](http://www.am.sanken.osaka-u.ac.jp/BiometricDB/GaitMVLP.html),
+excluding  identical-view cases.
+
+#### Fast
+With 8 NVIDIA 1080TI GPUs, it only takes **7 minutes** to conduct an evaluation on
+[OU-MVLP](http://www.am.sanken.osaka-u.ac.jp/BiometricDB/GaitMVLP.html) which contains 133,780 sequences
+and average 70 frames per sequence.
+
+## What's new
+The code and checkpoint for OUMVLP dataset have been released.
+See [OUMVLP](#oumvlp) for details.
+
+## Prerequisites
+
+- Python 3.6
+- PyTorch 0.4+
+- GPU
+
+
+## Getting started
+### Installation
+
+- (Not necessary) Install [Anaconda3](https://www.anaconda.com/download/)
+- Install [CUDA 9.0](https://developer.nvidia.com/cuda-90-download-archive)
+- install [cuDNN7.0](https://developer.nvidia.com/cudnn)
+- Install [PyTorch](http://pytorch.org/)
+
+Noted that our code is tested based on [PyTorch 0.4](http://pytorch.org/)
+
+### Dataset & Preparation
+Download [CASIA-B Dataset](http://www.cbsr.ia.ac.cn/english/Gait%20Databases.asp)
+
+**!!! ATTENTION !!! ATTENTION !!! ATTENTION !!!**
+
+Before training or test, please make sure you have prepared the dataset
+by this two steps:
+- **Step1:** Organize the directory as: 
+`your_dataset_path/subject_ids/walking_conditions/views`.
+E.g. `CASIA-B/001/nm-01/000/`.
+- **Step2:** Cut and align the raw silhouettes with `pretreatment.py`.
+(See [pretreatment](#pretreatment) for details.)
+Welcome to try different ways of pretreatment but note that
+the silhouettes after pretreatment **MUST have a size of 64x64**.
+
+Futhermore, you also can test our code on [OU-MVLP Dataset](http://www.am.sanken.osaka-u.ac.jp/BiometricDB/GaitMVLP.html).
+The number of channels and the training batchsize is slightly different for this dataset.
+For more detail, please refer to [our paper](https://arxiv.org/abs/1811.06186).
+
+#### Pretreatment
+`pretreatment.py` uses the alignment method in
+[this paper](https://ipsjcva.springeropen.com/articles/10.1186/s41074-018-0039-6).
+Pretreatment your dataset by
+```
+python pretreatment.py --input_path='root_path_of_raw_dataset' --output_path='root_path_for_output'
+```
+- `--input_path` **(NECESSARY)** Root path of raw dataset.
+- `--output_path` **(NECESSARY)** Root path for output.
+- `--log_file` Log file path. #Default: './pretreatment.log'
+- `--log` If set as True, all logs will be saved. 
+Otherwise, only warnings and errors will be saved. #Default: False
+- `--worker_num` How many subprocesses to use for data pretreatment. Default: 1
+
+### Configuration 
+
+In `config.py`, you might want to change the following settings:
+- `dataset_path` **(NECESSARY)** root path of the dataset 
+(for the above example, it is "gaitdata")
+- `WORK_PATH` path to save/load checkpoints
+- `CUDA_VISIBLE_DEVICES` indices of GPUs
+
+### Train
+Train a model by
+```bash
+python train.py
+```
+- `--cache` if set as TRUE all the training data will be loaded at once before the training start.
+This will accelerate the training.
+**Note that** if this arg is set as FALSE, samples will NOT be kept in the memory
+even they have been used in the former iterations. #Default: TRUE
+
+### Evaluation
+Evaluate the trained model by
+```bash
+python test.py
+```
+- `--iter` iteration of the checkpoint to load. #Default: 80000
+- `--batch_size` batch size of the parallel test. #Default: 1
+- `--cache` if set as TRUE all the test data will be loaded at once before the transforming start.
+This might accelerate the testing. #Default: FALSE
+
+It will output Rank@1 of all three walking conditions. 
+Note that the test is **parallelizable**. 
+To conduct a faster evaluation, you could use `--batch_size` to change the batch size for test.
+
+#### OUMVLP
+Since the huge differences between OUMVLP and CASIA-B, the network setting on OUMVLP is slightly different.
+- The alternated network's code can be found at `./work/OUMVLP_network`. Use them to replace the corresponding files in `./model/network`.
+- The checkpoint can be found [here](https://1drv.ms/u/s!AurT2TsSKdxQuWN8drzIv_phTR5m?e=Gfbl3m).
+- In `./config.py`, modify `'batch_size': (8, 16)` into `'batch_size': (32,16)`.
+- Prepare your OUMVLP dataset according to the instructions in [Dataset & Preparation](#dataset--preparation).
+
+## To Do List
+- Transformation: The script for transforming a set of silhouettes into a discriminative representation.
+
+## Authors & Contributors
+GaitSet is authored by
+[Hanqing Chao](https://www.linkedin.com/in/hanqing-chao-9aa42412b/), 
+[Yiwei He](https://www.linkedin.com/in/yiwei-he-4a6a6bbb/),
+[Junping Zhang](http://www.pami.fudan.edu.cn/~jpzhang/)
+and JianFeng Feng from Fudan Universiy.
+[Junping Zhang](http://www.pami.fudan.edu.cn/~jpzhang/)
+is the corresponding author.
+The code is developed by
+[Hanqing Chao](https://www.linkedin.com/in/hanqing-chao-9aa42412b/)
+and [Yiwei He](https://www.linkedin.com/in/yiwei-he-4a6a6bbb/).
+Currently, it is being maintained by
+[Hanqing Chao](https://www.linkedin.com/in/hanqing-chao-9aa42412b/)
+and Kun Wang.
+
+
+## Citation
+Please cite these papers in your publications if it helps your research:
+```
+@ARTICLE{chao2019gaitset,
+  author={Chao, Hanqing and Wang, Kun and He, Yiwei and Zhang, Junping and Feng, Jianfeng},
+  journal={IEEE Transactions on Pattern Analysis and Machine Intelligence}, 
+  title={GaitSet: Cross-view Gait Recognition through Utilizing Gait as a Deep Set}, 
+  year={2021},
+  pages={1-1},
+  doi={10.1109/TPAMI.2021.3057879}}
+```
+Link to paper:
+- [GaitSet: Cross-view Gait Recognition through Utilizing Gait as a Deep Set](https://ieeexplore.ieee.org/document/9351667)
+
+
+## License
+GaitSet is freely available for free non-commercial use, and may be redistributed under these conditions.
+For commercial queries, contact [Junping Zhang](http://www.pami.fudan.edu.cn/~jpzhang/).

config.py

@@ -0,0 +1,27 @@
+conf = {
+    "WORK_PATH": "./work",
+    "CUDA_VISIBLE_DEVICES": "0,1,2,3",
+    "data": {
+        'dataset_path': "your_dataset_path",
+        'resolution': '64',
+        'dataset': 'CASIA-B',
+        # In CASIA-B, data of subject #5 is incomplete.
+        # Thus, we ignore it in training.
+        # For more detail, please refer to
+        # function: utils.data_loader.load_data
+        'pid_num': 73,
+        'pid_shuffle': False,
+    },
+    "model": {
+        'hidden_dim': 256,
+        'lr': 1e-4,
+        'hard_or_full_trip': 'full',
+        'batch_size': (8, 16),
+        'restore_iter': 0,
+        'total_iter': 80000,
+        'margin': 0.2,
+        'num_workers': 3,
+        'frame_num': 30,
+        'model_name': 'GaitSet',
+    },
+}

gaitset/READ ME.md

@@ -0,0 +1,157 @@
+# GaitSet
+
+[![LICENSE](https://img.shields.io/badge/license-NPL%20(The%20996%20Prohibited%20License)-blue.svg)](https://github.com/996icu/996.ICU/blob/master/LICENSE)
+[![996.icu](https://img.shields.io/badge/link-996.icu-red.svg)](https://996.icu)
+
+GaitSet is a **flexible**, **effective** and **fast** network for cross-view gait recognition. The [paper](https://ieeexplore.ieee.org/document/9351667) has been published on IEEE TPAMI.
+
+#### Flexible 
+The input of GaitSet is a set of silhouettes. 
+
+- There are **NOT ANY constrains** on an input,
+which means it can contain **any number** of **non-consecutive** silhouettes filmed under **different viewpoints**
+with **different walking conditions**.
+
+- As the input is a set, the **permutation** of the elements in the input
+will **NOT change** the output at all.
+
+#### Effective
+It achieves **Rank@1=95.0%** on [CASIA-B](http://www.cbsr.ia.ac.cn/english/Gait%20Databases.asp) 
+and  **Rank@1=87.1%** on [OU-MVLP](http://www.am.sanken.osaka-u.ac.jp/BiometricDB/GaitMVLP.html),
+excluding  identical-view cases.
+
+#### Fast
+With 8 NVIDIA 1080TI GPUs, it only takes **7 minutes** to conduct an evaluation on
+[OU-MVLP](http://www.am.sanken.osaka-u.ac.jp/BiometricDB/GaitMVLP.html) which contains 133,780 sequences
+and average 70 frames per sequence.
+
+## What's new
+The code and checkpoint for OUMVLP dataset have been released.
+See [OUMVLP](#oumvlp) for details.
+
+## Prerequisites
+
+- Python 3.6
+- PyTorch 0.4+
+- GPU
+
+
+## Getting started
+### Installation
+
+- (Not necessary) Install [Anaconda3](https://www.anaconda.com/download/)
+- Install [CUDA 9.0](https://developer.nvidia.com/cuda-90-download-archive)
+- install [cuDNN7.0](https://developer.nvidia.com/cudnn)
+- Install [PyTorch](http://pytorch.org/)
+
+Noted that our code is tested based on [PyTorch 0.4](http://pytorch.org/)
+
+### Dataset & Preparation
+Download [CASIA-B Dataset](http://www.cbsr.ia.ac.cn/english/Gait%20Databases.asp)
+
+**!!! ATTENTION !!! ATTENTION !!! ATTENTION !!!**
+
+Before training or test, please make sure you have prepared the dataset
+by this two steps:
+- **Step1:** Organize the directory as: 
+`your_dataset_path/subject_ids/walking_conditions/views`.
+E.g. `CASIA-B/001/nm-01/000/`.
+- **Step2:** Cut and align the raw silhouettes with `pretreatment.py`.
+(See [pretreatment](#pretreatment) for details.)
+Welcome to try different ways of pretreatment but note that
+the silhouettes after pretreatment **MUST have a size of 64x64**.
+
+Futhermore, you also can test our code on [OU-MVLP Dataset](http://www.am.sanken.osaka-u.ac.jp/BiometricDB/GaitMVLP.html).
+The number of channels and the training batchsize is slightly different for this dataset.
+For more detail, please refer to [our paper](https://arxiv.org/abs/1811.06186).
+
+#### Pretreatment
+`pretreatment.py` uses the alignment method in
+[this paper](https://ipsjcva.springeropen.com/articles/10.1186/s41074-018-0039-6).
+Pretreatment your dataset by
+```
+python pretreatment.py --input_path='root_path_of_raw_dataset' --output_path='root_path_for_output'
+```
+- `--input_path` **(NECESSARY)** Root path of raw dataset.
+- `--output_path` **(NECESSARY)** Root path for output.
+- `--log_file` Log file path. #Default: './pretreatment.log'
+- `--log` If set as True, all logs will be saved. 
+Otherwise, only warnings and errors will be saved. #Default: False
+- `--worker_num` How many subprocesses to use for data pretreatment. Default: 1
+
+### Configuration 
+
+In `config.py`, you might want to change the following settings:
+- `dataset_path` **(NECESSARY)** root path of the dataset 
+(for the above example, it is "gaitdata")
+- `WORK_PATH` path to save/load checkpoints
+- `CUDA_VISIBLE_DEVICES` indices of GPUs
+
+### Train
+Train a model by
+```bash
+python train.py
+```
+- `--cache` if set as TRUE all the training data will be loaded at once before the training start.
+This will accelerate the training.
+**Note that** if this arg is set as FALSE, samples will NOT be kept in the memory
+even they have been used in the former iterations. #Default: TRUE
+
+### Evaluation
+Evaluate the trained model by
+```bash
+python test.py
+```
+- `--iter` iteration of the checkpoint to load. #Default: 80000
+- `--batch_size` batch size of the parallel test. #Default: 1
+- `--cache` if set as TRUE all the test data will be loaded at once before the transforming start.
+This might accelerate the testing. #Default: FALSE
+
+It will output Rank@1 of all three walking conditions. 
+Note that the test is **parallelizable**. 
+To conduct a faster evaluation, you could use `--batch_size` to change the batch size for test.
+
+#### OUMVLP
+Since the huge differences between OUMVLP and CASIA-B, the network setting on OUMVLP is slightly different.
+- The alternated network's code can be found at `./work/OUMVLP_network`. Use them to replace the corresponding files in `./model/network`.
+- The checkpoint can be found [here](https://1drv.ms/u/s!AurT2TsSKdxQuWN8drzIv_phTR5m?e=Gfbl3m).
+- In `./config.py`, modify `'batch_size': (8, 16)` into `'batch_size': (32,16)`.
+- Prepare your OUMVLP dataset according to the instructions in [Dataset & Preparation](#dataset--preparation).
+
+## To Do List
+- Transformation: The script for transforming a set of silhouettes into a discriminative representation.
+
+## Authors & Contributors
+GaitSet is authored by
+[Hanqing Chao](https://www.linkedin.com/in/hanqing-chao-9aa42412b/), 
+[Yiwei He](https://www.linkedin.com/in/yiwei-he-4a6a6bbb/),
+[Junping Zhang](http://www.pami.fudan.edu.cn/~jpzhang/)
+and JianFeng Feng from Fudan Universiy.
+[Junping Zhang](http://www.pami.fudan.edu.cn/~jpzhang/)
+is the corresponding author.
+The code is developed by
+[Hanqing Chao](https://www.linkedin.com/in/hanqing-chao-9aa42412b/)
+and [Yiwei He](https://www.linkedin.com/in/yiwei-he-4a6a6bbb/).
+Currently, it is being maintained by
+[Hanqing Chao](https://www.linkedin.com/in/hanqing-chao-9aa42412b/)
+and Kun Wang.
+
+
+## Citation
+Please cite these papers in your publications if it helps your research:
+```
+@ARTICLE{chao2019gaitset,
+  author={Chao, Hanqing and Wang, Kun and He, Yiwei and Zhang, Junping and Feng, Jianfeng},
+  journal={IEEE Transactions on Pattern Analysis and Machine Intelligence}, 
+  title={GaitSet: Cross-view Gait Recognition through Utilizing Gait as a Deep Set}, 
+  year={2021},
+  pages={1-1},
+  doi={10.1109/TPAMI.2021.3057879}}
+```
+Link to paper:
+- [GaitSet: Cross-view Gait Recognition through Utilizing Gait as a Deep Set](https://ieeexplore.ieee.org/document/9351667)
+
+
+## License
+GaitSet is freely available for free non-commercial use, and may be redistributed under these conditions.
+For commercial queries, contact [Junping Zhang](http://www.pami.fudan.edu.cn/~jpzhang/).

gaitset/basic_blocks.py

@@ -0,0 +1,49 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class BasicConv2d(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, **kwargs):
+        super(BasicConv2d, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, bias=False, **kwargs)
+
+    def forward(self, x):
+        x = self.conv(x)
+        return F.leaky_relu(x, inplace=True)
+
+
+class SetBlock(nn.Module):
+    def __init__(self, forward_block, pooling=False):
+        super(SetBlock, self).__init__()
+        self.forward_block = forward_block
+        self.pooling = pooling
+        if pooling:
+            self.pool2d = nn.MaxPool2d(2)
+    def forward(self, x):
+        n, s, c, h, w = x.size()
+        x = self.forward_block(x.view(-1,c,h,w))
+        if self.pooling:
+            x = self.pool2d(x)
+        _, c, h, w = x.size()
+        return x.view(n, s, c, h ,w)
+
+
+class HPM(nn.Module):
+    def __init__(self, in_dim, out_dim, bin_level_num=5):
+        super(HPM, self).__init__()
+        self.bin_num = [2**i for i in range(bin_level_num)]
+        self.fc_bin = nn.ParameterList([
+            nn.Parameter(
+                nn.init.xavier_uniform(
+                    torch.zeros(sum(self.bin_num), in_dim, out_dim)))])
+    def forward(self, x):
+        feature = list()
+        n, c, h, w = x.size()
+        for num_bin in self.bin_num:
+            z = x.view(n, c, num_bin, -1)
+            z = z.mean(3)+z.max(3)[0]
+            feature.append(z)
+        feature = torch.cat(feature, 2).permute(2, 0, 1).contiguous()
+        
+        feature = feature.matmul(self.fc_bin[0])
+        return feature.permute(1, 0, 2).contiguous()

gaitset/config.py

@@ -0,0 +1,27 @@
+conf = {
+    "WORK_PATH": "./work",
+    "CUDA_VISIBLE_DEVICES": "0,1,2,3",
+    "data": {
+        'dataset_path': "your_dataset_path",
+        'resolution': '64',
+        'dataset': 'CASIA-B',
+        # In CASIA-B, data of subject #5 is incomplete.
+        # Thus, we ignore it in training.
+        # For more detail, please refer to
+        # function: utils.data_loader.load_data
+        'pid_num': 73,
+        'pid_shuffle': False,
+    },
+    "model": {
+        'hidden_dim': 256,
+        'lr': 1e-4,
+        'hard_or_full_trip': 'full',
+        'batch_size': (8, 16),
+        'restore_iter': 0,
+        'total_iter': 80000,
+        'margin': 0.2,
+        'num_workers': 3,
+        'frame_num': 30,
+        'model_name': 'GaitSet',
+    },
+}

gaitset/model/_init_.py

@@ -0,0 +1,4 @@
+# -*- coding: utf-8 -*-
+# @Author  : admin
+# @Time    : 2018/11/16
+

gaitset/model/initialization.py

@@ -0,0 +1,58 @@
+# -*- coding: utf-8 -*-
+# @Author  : admin
+# @Time    : 2018/11/15
+import os
+from copy import deepcopy
+
+import numpy as np
+
+from .utils import load_data
+from .model import Model
+
+
+def initialize_data(config, train=False, test=False):
+    print("Initializing data source...")
+    train_source, test_source = load_data(**config['data'], cache=(train or test))
+    if train:
+        print("Loading training data...")
+        train_source.load_all_data()
+    if test:
+        print("Loading test data...")
+        test_source.load_all_data()
+    print("Data initialization complete.")
+    return train_source, test_source
+
+
+def initialize_model(config, train_source, test_source):
+    print("Initializing model...")
+    data_config = config['data']
+    model_config = config['model']
+    model_param = deepcopy(model_config)
+    model_param['train_source'] = train_source
+    model_param['test_source'] = test_source
+    model_param['train_pid_num'] = data_config['pid_num']
+    batch_size = int(np.prod(model_config['batch_size']))
+    model_param['save_name'] = '_'.join(map(str,[
+        model_config['model_name'],
+        data_config['dataset'],
+        data_config['pid_num'],
+        data_config['pid_shuffle'],
+        model_config['hidden_dim'],
+        model_config['margin'],
+        batch_size,
+        model_config['hard_or_full_trip'],
+        model_config['frame_num'],
+    ]))
+
+    m = Model(**model_param)
+    print("Model initialization complete.")
+    return m, model_param['save_name']
+
+
+def initialization(config, train=False, test=False):
+    print("Initialzing...")
+    WORK_PATH = config['WORK_PATH']
+    os.chdir(WORK_PATH)
+    os.environ["CUDA_VISIBLE_DEVICES"] = config["CUDA_VISIBLE_DEVICES"]
+    train_source, test_source = initialize_data(config, train, test)
+    return initialize_model(config, train_source, test_source)

gaitset/model/model.py

@@ -0,0 +1,272 @@
+import math
+import os
+import os.path as osp
+import random
+import sys
+from datetime import datetime
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.autograd as autograd
+import torch.optim as optim
+import torch.utils.data as tordata
+
+from .network import TripletLoss, SetNet
+from .utils import TripletSampler
+
+
+class Model:
+    def __init__(self,
+                 hidden_dim,
+                 lr,
+                 hard_or_full_trip,
+                 margin,
+                 num_workers,
+                 batch_size,
+                 restore_iter,
+                 total_iter,
+                 save_name,
+                 train_pid_num,
+                 frame_num,
+                 model_name,
+                 train_source,
+                 test_source,
+                 img_size=64):
+
+        self.save_name = save_name
+        self.train_pid_num = train_pid_num
+        self.train_source = train_source
+        self.test_source = test_source
+
+        self.hidden_dim = hidden_dim
+        self.lr = lr
+        self.hard_or_full_trip = hard_or_full_trip
+        self.margin = margin
+        self.frame_num = frame_num
+        self.num_workers = num_workers
+        self.batch_size = batch_size
+        self.model_name = model_name
+        self.P, self.M = batch_size
+
+        self.restore_iter = restore_iter
+        self.total_iter = total_iter
+
+        self.img_size = img_size
+
+        self.encoder = SetNet(self.hidden_dim).float()
+        self.encoder = nn.DataParallel(self.encoder)
+        self.triplet_loss = TripletLoss(self.P * self.M, self.hard_or_full_trip, self.margin).float()
+        self.triplet_loss = nn.DataParallel(self.triplet_loss)
+        self.encoder.cuda()
+        self.triplet_loss.cuda()
+
+        self.optimizer = optim.Adam([
+            {'params': self.encoder.parameters()},
+        ], lr=self.lr)
+
+        self.hard_loss_metric = []
+        self.full_loss_metric = []
+        self.full_loss_num = []
+        self.dist_list = []
+        self.mean_dist = 0.01
+
+        self.sample_type = 'all'
+
+    def collate_fn(self, batch):
+        batch_size = len(batch)
+        feature_num = len(batch[0][0])
+        seqs = [batch[i][0] for i in range(batch_size)]
+        frame_sets = [batch[i][1] for i in range(batch_size)]
+        view = [batch[i][2] for i in range(batch_size)]
+        seq_type = [batch[i][3] for i in range(batch_size)]
+        label = [batch[i][4] for i in range(batch_size)]
+        batch = [seqs, view, seq_type, label, None]
+
+        def select_frame(index):
+            sample = seqs[index]
+            frame_set = frame_sets[index]
+            if self.sample_type == 'random':
+                frame_id_list = random.choices(frame_set, k=self.frame_num)
+                _ = [feature.loc[frame_id_list].values for feature in sample]
+            else:
+                _ = [feature.values for feature in sample]
+            return _
+
+        seqs = list(map(select_frame, range(len(seqs))))
+
+        if self.sample_type == 'random':
+            seqs = [np.asarray([seqs[i][j] for i in range(batch_size)]) for j in range(feature_num)]
+        else:
+            gpu_num = min(torch.cuda.device_count(), batch_size)
+            batch_per_gpu = math.ceil(batch_size / gpu_num)
+            batch_frames = [[
+                                len(frame_sets[i])
+                                for i in range(batch_per_gpu * _, batch_per_gpu * (_ + 1))
+                                if i < batch_size
+                                ] for _ in range(gpu_num)]
+            if len(batch_frames[-1]) != batch_per_gpu:
+                for _ in range(batch_per_gpu - len(batch_frames[-1])):
+                    batch_frames[-1].append(0)
+            max_sum_frame = np.max([np.sum(batch_frames[_]) for _ in range(gpu_num)])
+            seqs = [[
+                        np.concatenate([
+                                           seqs[i][j]
+                                           for i in range(batch_per_gpu * _, batch_per_gpu * (_ + 1))
+                                           if i < batch_size
+                                           ], 0) for _ in range(gpu_num)]
+                    for j in range(feature_num)]
+            seqs = [np.asarray([
+                                   np.pad(seqs[j][_],
+                                          ((0, max_sum_frame - seqs[j][_].shape[0]), (0, 0), (0, 0)),
+                                          'constant',
+                                          constant_values=0)
+                                   for _ in range(gpu_num)])
+                    for j in range(feature_num)]
+            batch[4] = np.asarray(batch_frames)
+
+        batch[0] = seqs
+        return batch
+
+    def fit(self):
+        if self.restore_iter != 0:
+            self.load(self.restore_iter)
+
+        self.encoder.train()
+        self.sample_type = 'random'
+        for param_group in self.optimizer.param_groups:
+            param_group['lr'] = self.lr
+        triplet_sampler = TripletSampler(self.train_source, self.batch_size)
+        train_loader = tordata.DataLoader(
+            dataset=self.train_source,
+            batch_sampler=triplet_sampler,
+            collate_fn=self.collate_fn,
+            num_workers=self.num_workers)
+
+        train_label_set = list(self.train_source.label_set)
+        train_label_set.sort()
+
+        _time1 = datetime.now()
+        for seq, view, seq_type, label, batch_frame in train_loader:
+            self.restore_iter += 1
+            self.optimizer.zero_grad()
+
+            for i in range(len(seq)):
+                seq[i] = self.np2var(seq[i]).float()
+            if batch_frame is not None:
+                batch_frame = self.np2var(batch_frame).int()
+
+            feature, label_prob = self.encoder(*seq, batch_frame)
+
+            target_label = [train_label_set.index(l) for l in label]
+            target_label = self.np2var(np.array(target_label)).long()
+
+            triplet_feature = feature.permute(1, 0, 2).contiguous()
+            triplet_label = target_label.unsqueeze(0).repeat(triplet_feature.size(0), 1)
+            (full_loss_metric, hard_loss_metric, mean_dist, full_loss_num
+             ) = self.triplet_loss(triplet_feature, triplet_label)
+            if self.hard_or_full_trip == 'hard':
+                loss = hard_loss_metric.mean()
+            elif self.hard_or_full_trip == 'full':
+                loss = full_loss_metric.mean()
+
+            self.hard_loss_metric.append(hard_loss_metric.mean().data.cpu().numpy())
+            self.full_loss_metric.append(full_loss_metric.mean().data.cpu().numpy())
+            self.full_loss_num.append(full_loss_num.mean().data.cpu().numpy())
+            self.dist_list.append(mean_dist.mean().data.cpu().numpy())
+
+            if loss > 1e-9:
+                loss.backward()
+                self.optimizer.step()
+
+            if self.restore_iter % 1000 == 0:
+                print(datetime.now() - _time1)
+                _time1 = datetime.now()
+
+            if self.restore_iter % 100 == 0:
+                self.save()
+                print('iter {}:'.format(self.restore_iter), end='')
+                print(', hard_loss_metric={0:.8f}'.format(np.mean(self.hard_loss_metric)), end='')
+                print(', full_loss_metric={0:.8f}'.format(np.mean(self.full_loss_metric)), end='')
+                print(', full_loss_num={0:.8f}'.format(np.mean(self.full_loss_num)), end='')
+                self.mean_dist = np.mean(self.dist_list)
+                print(', mean_dist={0:.8f}'.format(self.mean_dist), end='')
+                print(', lr=%f' % self.optimizer.param_groups[0]['lr'], end='')
+                print(', hard or full=%r' % self.hard_or_full_trip)
+                sys.stdout.flush()
+                self.hard_loss_metric = []
+                self.full_loss_metric = []
+                self.full_loss_num = []
+                self.dist_list = []
+
+            # Visualization using t-SNE
+            # if self.restore_iter % 500 == 0:
+            #     pca = TSNE(2)
+            #     pca_feature = pca.fit_transform(feature.view(feature.size(0), -1).data.cpu().numpy())
+            #     for i in range(self.P):
+            #         plt.scatter(pca_feature[self.M * i:self.M * (i + 1), 0],
+            #                     pca_feature[self.M * i:self.M * (i + 1), 1], label=label[self.M * i])
+            #
+            #     plt.show()
+
+            if self.restore_iter == self.total_iter:
+                break
+
+    def ts2var(self, x):
+        return autograd.Variable(x).cuda()
+
+    def np2var(self, x):
+        return self.ts2var(torch.from_numpy(x))
+
+    def transform(self, flag, batch_size=1):
+        self.encoder.eval()
+        source = self.test_source if flag == 'test' else self.train_source
+        self.sample_type = 'all'
+        data_loader = tordata.DataLoader(
+            dataset=source,
+            batch_size=batch_size,
+            sampler=tordata.sampler.SequentialSampler(source),
+            collate_fn=self.collate_fn,
+            num_workers=self.num_workers)
+
+        feature_list = list()
+        view_list = list()
+        seq_type_list = list()
+        label_list = list()
+
+        for i, x in enumerate(data_loader):
+            seq, view, seq_type, label, batch_frame = x
+            for j in range(len(seq)):
+                seq[j] = self.np2var(seq[j]).float()
+            if batch_frame is not None:
+                batch_frame = self.np2var(batch_frame).int()
+            # print(batch_frame, np.sum(batch_frame))
+
+            feature, _ = self.encoder(*seq, batch_frame)
+            n, num_bin, _ = feature.size()
+            feature_list.append(feature.view(n, -1).data.cpu().numpy())
+            view_list += view
+            seq_type_list += seq_type
+            label_list += label
+
+        return np.concatenate(feature_list, 0), view_list, seq_type_list, label_list
+
+    def save(self):
+        os.makedirs(osp.join('checkpoint', self.model_name), exist_ok=True)
+        torch.save(self.encoder.state_dict(),
+                   osp.join('checkpoint', self.model_name,
+                            '{}-{:0>5}-encoder.ptm'.format(
+                                self.save_name, self.restore_iter)))
+        torch.save(self.optimizer.state_dict(),
+                   osp.join('checkpoint', self.model_name,
+                            '{}-{:0>5}-optimizer.ptm'.format(
+                                self.save_name, self.restore_iter)))
+
+    # restore_iter: iteration index of the checkpoint to load
+    def load(self, restore_iter):
+        self.encoder.load_state_dict(torch.load(osp.join(
+            'checkpoint', self.model_name,
+            '{}-{:0>5}-encoder.ptm'.format(self.save_name, restore_iter))))
+        self.optimizer.load_state_dict(torch.load(osp.join(
+            'checkpoint', self.model_name,
+            '{}-{:0>5}-optimizer.ptm'.format(self.save_name, restore_iter))))

gaitset/model/network/_init_.py

@@ -0,0 +1,2 @@
+from .triplet import TripletLoss
+from .gaitset import SetNet

gaitset/model/network/basic_blocks.py

@@ -0,0 +1,28 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class BasicConv2d(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, **kwargs):
+        super(BasicConv2d, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, bias=False, **kwargs)
+
+    def forward(self, x):
+        x = self.conv(x)
+        return F.leaky_relu(x, inplace=True)
+
+
+class SetBlock(nn.Module):
+    def __init__(self, forward_block, pooling=False):
+        super(SetBlock, self).__init__()
+        self.forward_block = forward_block
+        self.pooling = pooling
+        if pooling:
+            self.pool2d = nn.MaxPool2d(2)
+    def forward(self, x):
+        n, s, c, h, w = x.size()
+        x = self.forward_block(x.view(-1,c,h,w))
+        if self.pooling:
+            x = self.pool2d(x)
+        _, c, h, w = x.size()
+        return x.view(n, s, c, h ,w)

gaitset/model/network/gaitset.py

@@ -0,0 +1,120 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+from .basic_blocks import SetBlock, BasicConv2d
+
+
+class SetNet(nn.Module):
+    def __init__(self, hidden_dim):
+        super(SetNet, self).__init__()
+        self.hidden_dim = hidden_dim
+        self.batch_frame = None
+
+        _set_in_channels = 1
+        _set_channels = [32, 64, 128]
+        self.set_layer1 = SetBlock(BasicConv2d(_set_in_channels, _set_channels[0], 5, padding=2))
+        self.set_layer2 = SetBlock(BasicConv2d(_set_channels[0], _set_channels[0], 3, padding=1), True)
+        self.set_layer3 = SetBlock(BasicConv2d(_set_channels[0], _set_channels[1], 3, padding=1))
+        self.set_layer4 = SetBlock(BasicConv2d(_set_channels[1], _set_channels[1], 3, padding=1), True)
+        self.set_layer5 = SetBlock(BasicConv2d(_set_channels[1], _set_channels[2], 3, padding=1))
+        self.set_layer6 = SetBlock(BasicConv2d(_set_channels[2], _set_channels[2], 3, padding=1))
+
+        _gl_in_channels = 32
+        _gl_channels = [64, 128]
+        self.gl_layer1 = BasicConv2d(_gl_in_channels, _gl_channels[0], 3, padding=1)
+        self.gl_layer2 = BasicConv2d(_gl_channels[0], _gl_channels[0], 3, padding=1)
+        self.gl_layer3 = BasicConv2d(_gl_channels[0], _gl_channels[1], 3, padding=1)
+        self.gl_layer4 = BasicConv2d(_gl_channels[1], _gl_channels[1], 3, padding=1)
+        self.gl_pooling = nn.MaxPool2d(2)
+
+        self.bin_num = [1, 2, 4, 8, 16]
+        self.fc_bin = nn.ParameterList([
+            nn.Parameter(
+                nn.init.xavier_uniform_(
+                    torch.zeros(sum(self.bin_num) * 2, 128, hidden_dim)))])
+
+        for m in self.modules():
+            if isinstance(m, (nn.Conv2d, nn.Conv1d)):
+                nn.init.xavier_uniform_(m.weight.data)
+            elif isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight.data)
+                nn.init.constant(m.bias.data, 0.0)
+            elif isinstance(m, (nn.BatchNorm2d, nn.BatchNorm1d)):
+                nn.init.normal(m.weight.data, 1.0, 0.02)
+                nn.init.constant(m.bias.data, 0.0)
+
+    def frame_max(self, x):
+        if self.batch_frame is None:
+            return torch.max(x, 1)
+        else:
+            _tmp = [
+                torch.max(x[:, self.batch_frame[i]:self.batch_frame[i + 1], :, :, :], 1)
+                for i in range(len(self.batch_frame) - 1)
+                ]
+            max_list = torch.cat([_tmp[i][0] for i in range(len(_tmp))], 0)
+            arg_max_list = torch.cat([_tmp[i][1] for i in range(len(_tmp))], 0)
+            return max_list, arg_max_list
+
+    def frame_median(self, x):
+        if self.batch_frame is None:
+            return torch.median(x, 1)
+        else:
+            _tmp = [
+                torch.median(x[:, self.batch_frame[i]:self.batch_frame[i + 1], :, :, :], 1)
+                for i in range(len(self.batch_frame) - 1)
+                ]
+            median_list = torch.cat([_tmp[i][0] for i in range(len(_tmp))], 0)
+            arg_median_list = torch.cat([_tmp[i][1] for i in range(len(_tmp))], 0)
+            return median_list, arg_median_list
+
+    def forward(self, silho, batch_frame=None):
+        # n: batch_size, s: frame_num, k: keypoints_num, c: channel
+        if batch_frame is not None:
+            batch_frame = batch_frame[0].data.cpu().numpy().tolist()
+            _ = len(batch_frame)
+            for i in range(len(batch_frame)):
+                if batch_frame[-(i + 1)] != 0:
+                    break
+                else:
+                    _ -= 1
+            batch_frame = batch_frame[:_]
+            frame_sum = np.sum(batch_frame)
+            if frame_sum < silho.size(1):
+                silho = silho[:, :frame_sum, :, :]
+            self.batch_frame = [0] + np.cumsum(batch_frame).tolist()
+        n = silho.size(0)
+        x = silho.unsqueeze(2)
+        del silho
+
+        x = self.set_layer1(x)
+        x = self.set_layer2(x)
+        gl = self.gl_layer1(self.frame_max(x)[0])
+        gl = self.gl_layer2(gl)
+        gl = self.gl_pooling(gl)
+
+        x = self.set_layer3(x)
+        x = self.set_layer4(x)
+        gl = self.gl_layer3(gl + self.frame_max(x)[0])
+        gl = self.gl_layer4(gl)
+
+        x = self.set_layer5(x)
+        x = self.set_layer6(x)
+        x = self.frame_max(x)[0]
+        gl = gl + x
+
+        feature = list()
+        n, c, h, w = gl.size()
+        for num_bin in self.bin_num:
+            z = x.view(n, c, num_bin, -1)
+            z = z.mean(3) + z.max(3)[0]
+            feature.append(z)
+            z = gl.view(n, c, num_bin, -1)
+            z = z.mean(3) + z.max(3)[0]
+            feature.append(z)
+        feature = torch.cat(feature, 2).permute(2, 0, 1).contiguous()
+
+        feature = feature.matmul(self.fc_bin[0])
+        feature = feature.permute(1, 0, 2).contiguous()
+
+        return feature, None

gaitset/model/network/triplet.py

@@ -0,0 +1,45 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class TripletLoss(nn.Module):
+    def __init__(self, batch_size, hard_or_full, margin):
+        super(TripletLoss, self).__init__()
+        self.batch_size = batch_size
+        self.margin = margin
+
+    def forward(self, feature, label):
+        # feature: [n, m, d], label: [n, m]
+        n, m, d = feature.size()
+        hp_mask = (label.unsqueeze(1) == label.unsqueeze(2)).byte().view(-1)
+        hn_mask = (label.unsqueeze(1) != label.unsqueeze(2)).byte().view(-1)
+
+        dist = self.batch_dist(feature)
+        mean_dist = dist.mean(1).mean(1)
+        dist = dist.view(-1)
+        # hard
+        hard_hp_dist = torch.max(torch.masked_select(dist, hp_mask).view(n, m, -1), 2)[0]
+        hard_hn_dist = torch.min(torch.masked_select(dist, hn_mask).view(n, m, -1), 2)[0]
+        hard_loss_metric = F.relu(self.margin + hard_hp_dist - hard_hn_dist).view(n, -1)
+
+        hard_loss_metric_mean = torch.mean(hard_loss_metric, 1)
+
+        # non-zero full
+        full_hp_dist = torch.masked_select(dist, hp_mask).view(n, m, -1, 1)
+        full_hn_dist = torch.masked_select(dist, hn_mask).view(n, m, 1, -1)
+        full_loss_metric = F.relu(self.margin + full_hp_dist - full_hn_dist).view(n, -1)
+
+        full_loss_metric_sum = full_loss_metric.sum(1)
+        full_loss_num = (full_loss_metric != 0).sum(1).float()
+
+        full_loss_metric_mean = full_loss_metric_sum / full_loss_num
+        full_loss_metric_mean[full_loss_num == 0] = 0
+
+        return full_loss_metric_mean, hard_loss_metric_mean, mean_dist, full_loss_num
+
+    def batch_dist(self, x):
+        x2 = torch.sum(x ** 2, 2)
+        dist = x2.unsqueeze(2) + x2.unsqueeze(2).transpose(1, 2) - 2 * torch.matmul(x, x.transpose(1, 2))
+        dist = torch.sqrt(F.relu(dist))
+        return dist

gaitset/model/utils/_init_.py

@@ -0,0 +1,4 @@
+from .data_loader import load_data
+from .data_set import DataSet
+from .evaluator import evaluation
+from .sampler import TripletSampler

gaitset/model/utils/data_loader.py

@@ -0,0 +1,60 @@
+import os
+import os.path as osp
+
+import numpy as np
+
+from .data_set import DataSet
+
+
+def load_data(dataset_path, resolution, dataset, pid_num, pid_shuffle, cache=True):
+    seq_dir = list()
+    view = list()
+    seq_type = list()
+    label = list()
+
+    for _label in sorted(list(os.listdir(dataset_path))):
+        # In CASIA-B, data of subject #5 is incomplete.
+        # Thus, we ignore it in training.
+        if dataset == 'CASIA-B' and _label == '005':
+            continue
+        label_path = osp.join(dataset_path, _label)
+        for _seq_type in sorted(list(os.listdir(label_path))):
+            seq_type_path = osp.join(label_path, _seq_type)
+            for _view in sorted(list(os.listdir(seq_type_path))):
+                _seq_dir = osp.join(seq_type_path, _view)
+                seqs = os.listdir(_seq_dir)
+                if len(seqs) > 0:
+                    seq_dir.append([_seq_dir])
+                    label.append(_label)
+                    seq_type.append(_seq_type)
+                    view.append(_view)
+
+    pid_fname = osp.join('partition', '{}_{}_{}.npy'.format(
+        dataset, pid_num, pid_shuffle))
+    if not osp.exists(pid_fname):
+        pid_list = sorted(list(set(label)))
+        if pid_shuffle:
+            np.random.shuffle(pid_list)
+        pid_list = [pid_list[0:pid_num], pid_list[pid_num:]]
+        os.makedirs('partition', exist_ok=True)
+        np.save(pid_fname, pid_list)
+
+    pid_list = np.load(pid_fname)
+    train_list = pid_list[0]
+    test_list = pid_list[1]
+    train_source = DataSet(
+        [seq_dir[i] for i, l in enumerate(label) if l in train_list],
+        [label[i] for i, l in enumerate(label) if l in train_list],
+        [seq_type[i] for i, l in enumerate(label) if l in train_list],
+        [view[i] for i, l in enumerate(label)
+         if l in train_list],
+        cache, resolution)
+    test_source = DataSet(
+        [seq_dir[i] for i, l in enumerate(label) if l in test_list],
+        [label[i] for i, l in enumerate(label) if l in test_list],
+        [seq_type[i] for i, l in enumerate(label) if l in test_list],
+        [view[i] for i, l in enumerate(label)
+         if l in test_list],
+        cache, resolution)
+
+    return train_source, test_source

gaitset/model/utils/data_set.py

@@ -0,0 +1,90 @@
+import torch.utils.data as tordata
+import numpy as np
+import os.path as osp
+import os
+import pickle
+import cv2
+import xarray as xr
+
+
+class DataSet(tordata.Dataset):
+    def __init__(self, seq_dir, label, seq_type, view, cache, resolution):
+        self.seq_dir = seq_dir
+        self.view = view
+        self.seq_type = seq_type
+        self.label = label
+        self.cache = cache
+        self.resolution = int(resolution)
+        self.cut_padding = int(float(resolution)/64*10)
+        self.data_size = len(self.label)
+        self.data = [None] * self.data_size
+        self.frame_set = [None] * self.data_size
+
+        self.label_set = set(self.label)
+        self.seq_type_set = set(self.seq_type)
+        self.view_set = set(self.view)
+        _ = np.zeros((len(self.label_set),
+                      len(self.seq_type_set),
+                      len(self.view_set))).astype('int')
+        _ -= 1
+        self.index_dict = xr.DataArray(
+            _,
+            coords={'label': sorted(list(self.label_set)),
+                    'seq_type': sorted(list(self.seq_type_set)),
+                    'view': sorted(list(self.view_set))},
+            dims=['label', 'seq_type', 'view'])
+
+        for i in range(self.data_size):
+            _label = self.label[i]
+            _seq_type = self.seq_type[i]
+            _view = self.view[i]
+            self.index_dict.loc[_label, _seq_type, _view] = i
+
+    def load_all_data(self):
+        for i in range(self.data_size):
+            self.load_data(i)
+
+    def load_data(self, index):
+        return self.__getitem__(index)
+
+    def __loader__(self, path):
+        return self.img2xarray(
+            path)[:, :, self.cut_padding:-self.cut_padding].astype(
+            'float32') / 255.0
+
+    def __getitem__(self, index):
+        # pose sequence sampling
+        if not self.cache:
+            data = [self.__loader__(_path) for _path in self.seq_dir[index]]
+            frame_set = [set(feature.coords['frame'].values.tolist()) for feature in data]
+            frame_set = list(set.intersection(*frame_set))
+        elif self.data[index] is None:
+            data = [self.__loader__(_path) for _path in self.seq_dir[index]]
+            frame_set = [set(feature.coords['frame'].values.tolist()) for feature in data]
+            frame_set = list(set.intersection(*frame_set))
+            self.data[index] = data
+            self.frame_set[index] = frame_set
+        else:
+            data = self.data[index]
+            frame_set = self.frame_set[index]
+
+        return data, frame_set, self.view[
+            index], self.seq_type[index], self.label[index],
+
+    def img2xarray(self, flie_path):
+        imgs = sorted(list(os.listdir(flie_path)))
+        frame_list = [np.reshape(
+            cv2.imread(osp.join(flie_path, _img_path)),
+            [self.resolution, self.resolution, -1])[:, :, 0]
+                      for _img_path in imgs
+                      if osp.isfile(osp.join(flie_path, _img_path))]
+        num_list = list(range(len(frame_list)))
+        data_dict = xr.DataArray(
+            frame_list,
+            coords={'frame': num_list},
+            dims=['frame', 'img_y', 'img_x'],
+        )
+        return data_dict
+
+    def __len__(self):
+        return len(self.label)

gaitset/model/utils/evaluator.py

@@ -0,0 +1,49 @@
+import torch
+import torch.nn.functional as F
+import numpy as np
+
+
+def cuda_dist(x, y):
+    x = torch.from_numpy(x).cuda()
+    y = torch.from_numpy(y).cuda()
+    dist = torch.sum(x ** 2, 1).unsqueeze(1) + torch.sum(y ** 2, 1).unsqueeze(
+        1).transpose(0, 1) - 2 * torch.matmul(x, y.transpose(0, 1))
+    dist = torch.sqrt(F.relu(dist))
+    return dist
+
+
+def evaluation(data, config):
+    dataset = config['dataset'].split('-')[0]
+    feature, view, seq_type, label = data
+    label = np.array(label)
+    view_list = list(set(view))
+    view_list.sort()
+    view_num = len(view_list)
+    sample_num = len(feature)
+
+    probe_seq_dict = {'CASIA': [['nm-05', 'nm-06'], ['bg-01', 'bg-02'], ['cl-01', 'cl-02']],
+                      'OUMVLP': [['00']]}
+    gallery_seq_dict = {'CASIA': [['nm-01', 'nm-02', 'nm-03', 'nm-04']],
+                        'OUMVLP': [['01']]}
+
+    num_rank = 5
+    acc = np.zeros([len(probe_seq_dict[dataset]), view_num, view_num, num_rank])
+    for (p, probe_seq) in enumerate(probe_seq_dict[dataset]):
+        for gallery_seq in gallery_seq_dict[dataset]:
+            for (v1, probe_view) in enumerate(view_list):
+                for (v2, gallery_view) in enumerate(view_list):
+                    gseq_mask = np.isin(seq_type, gallery_seq) & np.isin(view, [gallery_view])
+                    gallery_x = feature[gseq_mask, :]
+                    gallery_y = label[gseq_mask]
+
+                    pseq_mask = np.isin(seq_type, probe_seq) & np.isin(view, [probe_view])
+                    probe_x = feature[pseq_mask, :]
+                    probe_y = label[pseq_mask]
+
+                    dist = cuda_dist(probe_x, gallery_x)
+                    idx = dist.sort(1)[1].cpu().numpy()
+                    acc[p, v1, v2, :] = np.round(
+                        np.sum(np.cumsum(np.reshape(probe_y, [-1, 1]) == gallery_y[idx[:, 0:num_rank]], 1) > 0,
+                               0) * 100 / dist.shape[0], 2)
+
+    return acc

gaitset/model/utils/sampler.py

@@ -0,0 +1,26 @@
+import torch.utils.data as tordata
+import random
+
+
+class TripletSampler(tordata.sampler.Sampler):
+    def __init__(self, dataset, batch_size):
+        self.dataset = dataset
+        self.batch_size = batch_size
+
+    def __iter__(self):
+        while (True):
+            sample_indices = list()
+            pid_list = random.sample(
+                list(self.dataset.label_set),
+                self.batch_size[0])
+            for pid in pid_list:
+                _index = self.dataset.index_dict.loc[pid, :, :].values
+                _index = _index[_index > 0].flatten().tolist()
+                _index = random.choices(
+                    _index,
+                    k=self.batch_size[1])
+                sample_indices += _index
+            yield sample_indices
+
+    def __len__(self):
+        return self.dataset.data_size

gaitset/pretreatment.py

@@ -0,0 +1,192 @@
+
+import os
+from scipy import misc as scisc
+import cv2
+import numpy as np
+from warnings import warn
+from time import sleep
+import argparse
+
+from multiprocessing import Pool
+from multiprocessing import TimeoutError as MP_TimeoutError
+
+START = "START"
+FINISH = "FINISH"
+WARNING = "WARNING"
+FAIL = "FAIL"
+
+
+def boolean_string(s):
+    if s.upper() not in {'FALSE', 'TRUE'}:
+        raise ValueError('Not a valid boolean string')
+    return s.upper() == 'TRUE'
+
+
+parser = argparse.ArgumentParser(description='Test')
+parser.add_argument('--input_path', default='', type=str,
+                    help='Root path of raw dataset.')
+parser.add_argument('--output_path', default='', type=str,
+                    help='Root path for output.')
+parser.add_argument('--log_file', default='./pretreatment.log', type=str,
+                    help='Log file path. Default: ./pretreatment.log')
+parser.add_argument('--log', default=False, type=boolean_string,
+                    help='If set as True, all logs will be saved. '
+                         'Otherwise, only warnings and errors will be saved.'
+                         'Default: False')
+parser.add_argument('--worker_num', default=1, type=int,
+                    help='How many subprocesses to use for data pretreatment. '
+                         'Default: 1')
+opt = parser.parse_args()
+
+INPUT_PATH = opt.input_path
+OUTPUT_PATH = opt.output_path
+IF_LOG = opt.log
+LOG_PATH = opt.log_file
+WORKERS = opt.worker_num
+
+T_H = 64
+T_W = 64
+
+
+def log2str(pid, comment, logs):
+    str_log = ''
+    if type(logs) is str:
+        logs = [logs]
+    for log in logs:
+        str_log += "# JOB %d : --%s-- %s\n" % (
+            pid, comment, log)
+    return str_log
+
+
+def log_print(pid, comment, logs):
+    str_log = log2str(pid, comment, logs)
+    if comment in [WARNING, FAIL]:
+        with open(LOG_PATH, 'a') as log_f:
+            log_f.write(str_log)
+    if comment in [START, FINISH]:
+        if pid % 500 != 0:
+            return
+    print(str_log, end='')
+
+
+def cut_img(img, seq_info, frame_name, pid):
+    # A silhouette contains too little white pixels
+    # might be not valid for identification.
+    if img.sum() <= 10000:
+        message = 'seq:%s, frame:%s, no data, %d.' % (
+            '-'.join(seq_info), frame_name, img.sum())
+        warn(message)
+        log_print(pid, WARNING, message)
+        return None
+    # Get the top and bottom point
+    y = img.sum(axis=1)
+    y_top = (y != 0).argmax(axis=0)
+    y_btm = (y != 0).cumsum(axis=0).argmax(axis=0)
+    img = img[y_top:y_btm + 1, :]
+    # As the height of a person is larger than the width,
+    # use the height to calculate resize ratio.
+    _r = img.shape[1] / img.shape[0]
+    _t_w = int(T_H * _r)
+    img = cv2.resize(img, (_t_w, T_H), interpolation=cv2.INTER_CUBIC)
+    # Get the median of x axis and regard it as the x center of the person.
+    sum_point = img.sum()
+    sum_column = img.sum(axis=0).cumsum()
+    x_center = -1
+    for i in range(sum_column.size):
+        if sum_column[i] > sum_point / 2:
+            x_center = i
+            break
+    if x_center < 0:
+        message = 'seq:%s, frame:%s, no center.' % (
+            '-'.join(seq_info), frame_name)
+        warn(message)
+        log_print(pid, WARNING, message)
+        return None
+    h_T_W = int(T_W / 2)
+    left = x_center - h_T_W
+    right = x_center + h_T_W
+    if left <= 0 or right >= img.shape[1]:
+        left += h_T_W
+        right += h_T_W
+        _ = np.zeros((img.shape[0], h_T_W))
+        img = np.concatenate([_, img, _], axis=1)
+    img = img[:, left:right]
+    return img.astype('uint8')
+
+
+def cut_pickle(seq_info, pid):
+    seq_name = '-'.join(seq_info)
+    log_print(pid, START, seq_name)
+    seq_path = os.path.join(INPUT_PATH, *seq_info)
+    out_dir = os.path.join(OUTPUT_PATH, *seq_info)
+    frame_list = os.listdir(seq_path)
+    frame_list.sort()
+    count_frame = 0
+    for _frame_name in frame_list:
+        frame_path = os.path.join(seq_path, _frame_name)
+        img = cv2.imread(frame_path)[:, :, 0]
+        img = cut_img(img, seq_info, _frame_name, pid)
+        if img is not None:
+            # Save the cut img
+            save_path = os.path.join(out_dir, _frame_name)
+            scisc.imsave(save_path, img)
+            count_frame += 1
+    # Warn if the sequence contains less than 5 frames
+    if count_frame < 5:
+        message = 'seq:%s, less than 5 valid data.' % (
+            '-'.join(seq_info))
+        warn(message)
+        log_print(pid, WARNING, message)
+
+    log_print(pid, FINISH,
+              'Contain %d valid frames. Saved to %s.'
+              % (count_frame, out_dir))
+
+
+pool = Pool(WORKERS)
+results = list()
+pid = 0
+
+print('Pretreatment Start.\n'
+      'Input path: %s\n'
+      'Output path: %s\n'
+      'Log file: %s\n'
+      'Worker num: %d' % (
+          INPUT_PATH, OUTPUT_PATH, LOG_PATH, WORKERS))
+
+id_list = os.listdir(INPUT_PATH)
+id_list.sort()
+# Walk the input path
+for _id in id_list:
+    seq_type = os.listdir(os.path.join(INPUT_PATH, _id))
+    seq_type.sort()
+    for _seq_type in seq_type:
+        view = os.listdir(os.path.join(INPUT_PATH, _id, _seq_type))
+        view.sort()
+        for _view in view:
+            seq_info = [_id, _seq_type, _view]
+            out_dir = os.path.join(OUTPUT_PATH, *seq_info)
+            os.makedirs(out_dir)
+            results.append(
+                pool.apply_async(
+                    cut_pickle,
+                    args=(seq_info, pid)))
+            sleep(0.02)
+            pid += 1
+
+pool.close()
+unfinish = 1
+while unfinish > 0:
+    unfinish = 0
+    for i, res in enumerate(results):
+        try:
+            res.get(timeout=0.1)
+        except Exception as e:
+            if type(e) == MP_TimeoutError:
+                unfinish += 1
+                continue
+            else:
+                print('\n\n\nERROR OCCUR: PID ##%d##, ERRORTYPE: %s\n\n\n',
+                      i, type(e))
+                raise e
+pool.join()

gaitset/test.py

@@ -0,0 +1,80 @@
+from datetime import datetime
+import numpy as np
+import argparse
+
+from model.initialization import initialization
+from model.utils import evaluation
+from config import conf
+
+
+def boolean_string(s):
+    if s.upper() not in {'FALSE', 'TRUE'}:
+        raise ValueError('Not a valid boolean string')
+    return s.upper() == 'TRUE'
+
+
+parser = argparse.ArgumentParser(description='Test')
+parser.add_argument('--iter', default='80000', type=int,
+                    help='iter: iteration of the checkpoint to load. Default: 80000')
+parser.add_argument('--batch_size', default='1', type=int,
+                    help='batch_size: batch size for parallel test. Default: 1')
+parser.add_argument('--cache', default=False, type=boolean_string,
+                    help='cache: if set as TRUE all the test data will be loaded at once'
+                         ' before the transforming start. Default: FALSE')
+opt = parser.parse_args()
+
+
+# Exclude identical-view cases
+def de_diag(acc, each_angle=False):
+    result = np.sum(acc - np.diag(np.diag(acc)), 1) / 10.0
+    if not each_angle:
+        result = np.mean(result)
+    return result
+
+
+m = initialization(conf, test=opt.cache)[0]
+
+# load model checkpoint of iteration opt.iter
+print('Loading the model of iteration %d...' % opt.iter)
+m.load(opt.iter)
+print('Transforming...')
+time = datetime.now()
+test = m.transform('test', opt.batch_size)
+print('Evaluating...')
+acc = evaluation(test, conf['data'])
+print('Evaluation complete. Cost:', datetime.now() - time)
+
+# Print rank-1 accuracy of the best model
+# e.g.
+# ===Rank-1 (Include identical-view cases)===
+# NM: 95.405,     BG: 88.284,     CL: 72.041
+for i in range(1):
+    print('===Rank-%d (Include identical-view cases)===' % (i + 1))
+    print('NM: %.3f,\tBG: %.3f,\tCL: %.3f' % (
+        np.mean(acc[0, :, :, i]),
+        np.mean(acc[1, :, :, i]),
+        np.mean(acc[2, :, :, i])))
+
+# Print rank-1 accuracy of the best model，excluding identical-view cases
+# e.g.
+# ===Rank-1 (Exclude identical-view cases)===
+# NM: 94.964,     BG: 87.239,     CL: 70.355
+for i in range(1):
+    print('===Rank-%d (Exclude identical-view cases)===' % (i + 1))
+    print('NM: %.3f,\tBG: %.3f,\tCL: %.3f' % (
+        de_diag(acc[0, :, :, i]),
+        de_diag(acc[1, :, :, i]),
+        de_diag(acc[2, :, :, i])))
+
+# Print rank-1 accuracy of the best model (Each Angle)
+# e.g.
+# ===Rank-1 of each angle (Exclude identical-view cases)===
+# NM: [90.80 97.90 99.40 96.90 93.60 91.70 95.00 97.80 98.90 96.80 85.80]
+# BG: [83.80 91.20 91.80 88.79 83.30 81.00 84.10 90.00 92.20 94.45 79.00]
+# CL: [61.40 75.40 80.70 77.30 72.10 70.10 71.50 73.50 73.50 68.40 50.00]
+np.set_printoptions(precision=2, floatmode='fixed')
+for i in range(1):
+    print('===Rank-%d of each angle (Exclude identical-view cases)===' % (i + 1))
+    print('NM:', de_diag(acc[0, :, :, i], True))
+    print('BG:', de_diag(acc[1, :, :, i], True))
+    print('CL:', de_diag(acc[2, :, :, i], True))

gaitset/train.py

@@ -0,0 +1,22 @@
+from model.initialization import initialization
+from config import conf
+import argparse
+
+
+def boolean_string(s):
+    if s.upper() not in {'FALSE', 'TRUE'}:
+        raise ValueError('Not a valid boolean string')
+    return s.upper() == 'TRUE'
+
+
+parser = argparse.ArgumentParser(description='Train')
+parser.add_argument('--cache', default=True, type=boolean_string,
+                    help='cache: if set as TRUE all the training data will be loaded at once'
+                         ' before the training start. Default: TRUE')
+opt = parser.parse_args()
+
+m = initialization(conf, train=opt.cache)[0]
+
+print("Training START")
+m.fit()
+print("Training COMPLETE")

gaitset/work/basic_blocks.py

@@ -0,0 +1,49 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class BasicConv2d(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, **kwargs):
+        super(BasicConv2d, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, bias=False, **kwargs)
+
+    def forward(self, x):
+        x = self.conv(x)
+        return F.leaky_relu(x, inplace=True)
+
+
+class SetBlock(nn.Module):
+    def __init__(self, forward_block, pooling=False):
+        super(SetBlock, self).__init__()
+        self.forward_block = forward_block
+        self.pooling = pooling
+        if pooling:
+            self.pool2d = nn.MaxPool2d(2)
+    def forward(self, x):
+        n, s, c, h, w = x.size()
+        x = self.forward_block(x.view(-1,c,h,w))
+        if self.pooling:
+            x = self.pool2d(x)
+        _, c, h, w = x.size()
+        return x.view(n, s, c, h ,w)
+
+
+class HPM(nn.Module):
+    def __init__(self, in_dim, out_dim, bin_level_num=5):
+        super(HPM, self).__init__()
+        self.bin_num = [2**i for i in range(bin_level_num)]
+        self.fc_bin = nn.ParameterList([
+            nn.Parameter(
+                nn.init.xavier_uniform(
+                    torch.zeros(sum(self.bin_num), in_dim, out_dim)))])
+    def forward(self, x):
+        feature = list()
+        n, c, h, w = x.size()
+        for num_bin in self.bin_num:
+            z = x.view(n, c, num_bin, -1)
+            z = z.mean(3)+z.max(3)[0]
+            feature.append(z)
+        feature = torch.cat(feature, 2).permute(2, 0, 1).contiguous()
+        
+        feature = feature.matmul(self.fc_bin[0])
+        return feature.permute(1, 0, 2).contiguous()

gaitset/work/gaitset.py

@@ -0,0 +1,88 @@
+class SetNet(nn.Module):
+    def __init__(self, hidden_dim):
+        super(SetNet, self).__init__()
+        self.hidden_dim = hidden_dim
+        self.batch_frame = None
+        
+        _in_channels = 1
+        _channels = [64,128,256]
+        self.set_layer1 = SetBlock(BasicConv2d(_in_channels, _channels[0], 5, padding=2))
+        self.set_layer2 = SetBlock(BasicConv2d(_channels[0], _channels[0], 3, padding=1), True)
+        self.set_layer3 = SetBlock(BasicConv2d(_channels[0], _channels[1], 3, padding=1))
+        self.set_layer4 = SetBlock(BasicConv2d(_channels[1], _channels[1], 3, padding=1), True)
+        self.set_layer5 = SetBlock(BasicConv2d(_channels[1], _channels[2], 3, padding=1))
+        self.set_layer6 = SetBlock(BasicConv2d(_channels[2], _channels[2], 3, padding=1))
+        
+        self.gl_layer1 = BasicConv2d(_channels[0], _channels[1], 3, padding=1)
+        self.gl_layer2 = BasicConv2d(_channels[1], _channels[1], 3, padding=1)
+        self.gl_layer3 = BasicConv2d(_channels[1], _channels[2], 3, padding=1)
+        self.gl_layer4 = BasicConv2d(_channels[2], _channels[2], 3, padding=1)
+        self.gl_pooling = nn.MaxPool2d(2)
+        
+        self.gl_hpm = HPM(_channels[-1], hidden_dim)
+        self.x_hpm = HPM(_channels[-1], hidden_dim)
+        
+        for m in self.modules():
+            if isinstance(m, (nn.Conv2d, nn.Conv1d)):
+                nn.init.xavier_uniform(m.weight.data)
+            elif isinstance(m, nn.Linear):
+                nn.init.xavier_uniform(m.weight.data)
+                nn.init.constant(m.bias.data, 0.0)
+            elif isinstance(m, (nn.BatchNorm2d, nn.BatchNorm1d)):
+                nn.init.normal(m.weight.data, 1.0, 0.02)
+                nn.init.constant(m.bias.data, 0.0)
+                    
+
+    def frame_max(self, x):
+        if self.batch_frame is None:
+            return torch.max(x, 1)
+        else:
+            _tmp = [
+                torch.max(x[:, self.batch_frame[i]:self.batch_frame[i+1], :, :, :], 1)
+                for i in range(len(self.batch_frame)-1)
+            ]
+            max_list = torch.cat([_tmp[i][0] for i in range(len(_tmp))], 0)
+            arg_max_list = torch.cat([_tmp[i][1] for i in range(len(_tmp))], 0)
+            return max_list, arg_max_list
+
+                
+    def forward(self, silho, batch_frame=None):
+        silho = silho/255
+        # n: batch_size, s: frame_num, k: keypoints_num, c: channel
+        if batch_frame is not None:
+            batch_frame = batch_frame[0].data.cpu().numpy().tolist()
+            _ = len(batch_frame)
+            for i in range(len(batch_frame)):
+                if batch_frame[-(i+1)] != 0:
+                    break
+                else:
+                    _ -= 1
+            batch_frame = batch_frame[:_]
+            frame_sum = np.sum(batch_frame)
+            if frame_sum < silho.size(1):
+                silho = silho[:, :frame_sum,:,:]
+            self.batch_frame = [0]+np.cumsum(batch_frame).tolist()
+        n = silho.size(0)
+        x = silho.unsqueeze(2)
+        del silho
+        
+        x = self.set_layer1(x)
+        x = self.set_layer2(x)
+        gl = self.gl_layer1(self.frame_max(x)[0])
+        gl = self.gl_layer2(gl)
+        gl = self.gl_pooling(gl)
+        
+        x = self.set_layer3(x)
+        x = self.set_layer4(x)
+        gl = self.gl_layer3(gl+self.frame_max(x)[0])
+        gl = self.gl_layer4(gl)
+        
+        x = self.set_layer5(x)
+        x = self.set_layer6(x)
+        x = self.frame_max(x)[0]
+        gl = gl+x
+        
+        gl_f = self.gl_hpm(gl)
+        x_f = self.x_hpm(x)
+        
+        return torch.cat([gl_f, x_f], 1), None

test.py

@@ -0,0 +1,80 @@
+from datetime import datetime
+import numpy as np
+import argparse
+
+from model.initialization import initialization
+from model.utils import evaluation
+from config import conf
+
+
+def boolean_string(s):
+    if s.upper() not in {'FALSE', 'TRUE'}:
+        raise ValueError('Not a valid boolean string')
+    return s.upper() == 'TRUE'
+
+
+parser = argparse.ArgumentParser(description='Test')
+parser.add_argument('--iter', default='80000', type=int,
+                    help='iter: iteration of the checkpoint to load. Default: 80000')
+parser.add_argument('--batch_size', default='1', type=int,
+                    help='batch_size: batch size for parallel test. Default: 1')
+parser.add_argument('--cache', default=False, type=boolean_string,
+                    help='cache: if set as TRUE all the test data will be loaded at once'
+                         ' before the transforming start. Default: FALSE')
+opt = parser.parse_args()
+
+
+# Exclude identical-view cases
+def de_diag(acc, each_angle=False):
+    result = np.sum(acc - np.diag(np.diag(acc)), 1) / 10.0
+    if not each_angle:
+        result = np.mean(result)
+    return result
+
+
+m = initialization(conf, test=opt.cache)[0]
+
+# load model checkpoint of iteration opt.iter
+print('Loading the model of iteration %d...' % opt.iter)
+m.load(opt.iter)
+print('Transforming...')
+time = datetime.now()
+test = m.transform('test', opt.batch_size)
+print('Evaluating...')
+acc = evaluation(test, conf['data'])
+print('Evaluation complete. Cost:', datetime.now() - time)
+
+# Print rank-1 accuracy of the best model
+# e.g.
+# ===Rank-1 (Include identical-view cases)===
+# NM: 95.405,     BG: 88.284,     CL: 72.041
+for i in range(1):
+    print('===Rank-%d (Include identical-view cases)===' % (i + 1))
+    print('NM: %.3f,\tBG: %.3f,\tCL: %.3f' % (
+        np.mean(acc[0, :, :, i]),
+        np.mean(acc[1, :, :, i]),
+        np.mean(acc[2, :, :, i])))
+
+# Print rank-1 accuracy of the best model，excluding identical-view cases
+# e.g.
+# ===Rank-1 (Exclude identical-view cases)===
+# NM: 94.964,     BG: 87.239,     CL: 70.355
+for i in range(1):
+    print('===Rank-%d (Exclude identical-view cases)===' % (i + 1))
+    print('NM: %.3f,\tBG: %.3f,\tCL: %.3f' % (
+        de_diag(acc[0, :, :, i]),
+        de_diag(acc[1, :, :, i]),
+        de_diag(acc[2, :, :, i])))
+
+# Print rank-1 accuracy of the best model (Each Angle)
+# e.g.
+# ===Rank-1 of each angle (Exclude identical-view cases)===
+# NM: [90.80 97.90 99.40 96.90 93.60 91.70 95.00 97.80 98.90 96.80 85.80]
+# BG: [83.80 91.20 91.80 88.79 83.30 81.00 84.10 90.00 92.20 94.45 79.00]
+# CL: [61.40 75.40 80.70 77.30 72.10 70.10 71.50 73.50 73.50 68.40 50.00]
+np.set_printoptions(precision=2, floatmode='fixed')
+for i in range(1):
+    print('===Rank-%d of each angle (Exclude identical-view cases)===' % (i + 1))
+    print('NM:', de_diag(acc[0, :, :, i], True))
+    print('BG:', de_diag(acc[1, :, :, i], True))
+    print('CL:', de_diag(acc[2, :, :, i], True))

train.py

@@ -0,0 +1,22 @@
+from model.initialization import initialization
+from config import conf
+import argparse
+
+
+def boolean_string(s):
+    if s.upper() not in {'FALSE', 'TRUE'}:
+        raise ValueError('Not a valid boolean string')
+    return s.upper() == 'TRUE'
+
+
+parser = argparse.ArgumentParser(description='Train')
+parser.add_argument('--cache', default=True, type=boolean_string,
+                    help='cache: if set as TRUE all the training data will be loaded at once'
+                         ' before the training start. Default: TRUE')
+opt = parser.parse_args()
+
+m = initialization(conf, train=opt.cache)[0]
+
+print("Training START")
+m.fit()
+print("Training COMPLETE")