Merge pull request #37 from yizhou-wang/v1.1-mnet-tdc

V1.1 add mnet tdc modules

.gitignore

 __pycache__/
 *.egg-info/
+*.so
 
+build/
 data/
 checkpoints/
 results/

README.md

-# RODNet: Radar Object Detection using Cross-Modal Supervision
+# RODNet: Radar Object Detection Network
 
-This is the official implementation of our RODNet paper at WACV 2021. 
+This is the official implementation of our RODNet papers 
+at [WACV 2021](https://openaccess.thecvf.com/content/WACV2021/html/Wang_RODNet_Radar_Object_Detection_Using_Cross-Modal_Supervision_WACV_2021_paper.html) 
+and [IEEE J-STSP 2021](https://ieeexplore.ieee.org/abstract/document/9353210). 
 
-[[Paper]](https://openaccess.thecvf.com/content/WACV2021/html/Wang_RODNet_Radar_Object_Detection_Using_Cross-Modal_Supervision_WACV_2021_paper.html)
+[[Arxiv]](https://arxiv.org/abs/2102.05150)
 [[Dataset]](https://www.cruwdataset.org)
+[[ROD2021 Challenge]](https://codalab.lisn.upsaclay.fr/competitions/1063)
+[[Presentation]](https://youtu.be/UZbxI4o2-7g)
+[[Demo]](https://youtu.be/09HaDySa29I)
 
 ![RODNet Overview](./assets/images/overview.jpg?raw=true)
 
-Please cite our WACV 2021 paper if this repository is helpful for your research:
+Please cite our paper if this repository is helpful for your research:
 ```
 @inproceedings{wang2021rodnet,
   author={Wang, Yizhou and Jiang, Zhongyu and Gao, Xiangyu and Hwang, Jenq-Neng and Xing, Guanbin and Liu, Hui},
@@ -18,22 +23,40 @@ Please cite our WACV 2021 paper if this repository is helpful for your research:
   pages={504-513}
 }
 ```
+```
+@article{wang2021rodnet,
+  title={RODNet: A Real-Time Radar Object Detection Network Cross-Supervised by Camera-Radar Fused Object 3D Localization},
+  author={Wang, Yizhou and Jiang, Zhongyu and Li, Yudong and Hwang, Jenq-Neng and Xing, Guanbin and Liu, Hui},
+  journal={IEEE Journal of Selected Topics in Signal Processing},
+  volume={15},
+  number={4},
+  pages={954--967},
+  year={2021},
+  publisher={IEEE}
+}
+```
+
 ## Installation
 
 Create a conda environment for RODNet. Tested under Python 3.6, 3.7, 3.8.
-```
+```commandline
 conda create -n rodnet python=3.* -y
 conda activate rodnet
 ```
 
 Install pytorch.
-```
-conda install pytorch torchvision cudatoolkit=10.1 -c pytorch
+**Note:** If you are using Temporal Deformable Convolution (TDC), we only tested under `pytorch<=1.4` and `CUDA=10.1`. 
+Without TDC, you should be able to choose the latest versions. 
+If you met some issues with environment, feel free to raise an issue.
+```commandline
+conda install pytorch=1.4 torchvision cudatoolkit=10.1 -c pytorch  # if using TDC
+# OR
+conda install pytorch torchvision cudatoolkit=10.1 -c pytorch  # if not using TDC
 ```
 
 Install `cruw-devkit` package. 
 Please refer to [`cruw-devit`](https://github.com/yizhou-wang/cruw-devkit) repository for detailed instructions.
-```
+```commandline
 git clone https://github.com/yizhou-wang/cruw-devkit.git
 cd cruw-devkit
 pip install -e .
@@ -41,13 +64,15 @@ cd ..
 ```
 
 Setup RODNet package.
-```
+```commandline
 pip install -e .
 ```
+**Note:** If you are not using TDC, you can rename script `setup_wo_tdc.py` as `setup.py`, and run the above command. 
+This should allow you to use the latest cuda and pytorch version. 
 
 ## Prepare data for RODNet
 
-```
+```commandline
 python tools/prepare_dataset/prepare_data.py \
         --config configs/<CONFIG_FILE> \
         --data_root <DATASET_ROOT> \
@@ -57,7 +82,7 @@ python tools/prepare_dataset/prepare_data.py \
 
 ## Train models
 
-```
+```commandline
 python tools/train.py --config configs/<CONFIG_FILE> \
         --data_dir data/<DATA_FOLDER_NAME> \
         --log_dir checkpoints/
@@ -65,7 +90,7 @@ python tools/train.py --config configs/<CONFIG_FILE> \
 
 ## Inference
 
-```
+```commandline
 python tools/test.py --config configs/<CONFIG_FILE> \
         --data_dir data/<DATA_FOLDER_NAME> \
         --checkpoint <CHECKPOINT_PATH> \

configs/config_rodnet_v2_cdc_win16.py

+dataset_cfg = dict(
+    dataset_name='ROD2021',
+    base_root="/mnt/disk2/CRUW/ROD2021",
+    data_root="/mnt/disk2/CRUW/ROD2021/sequences",
+    anno_root="/mnt/disk2/CRUW/ROD2021/annotations",
+    anno_ext='.txt',
+    train=dict(
+        subdir='train',
+        # seqs=[],  # can choose from the subdir folder
+    ),
+    valid=dict(
+        subdir='valid',
+        seqs=[],
+    ),
+    test=dict(
+        subdir='test',
+        # seqs=[],  # can choose from the subdir folder
+    ),
+    demo=dict(
+        subdir='demo',
+        seqs=[],
+    ),
+)
+
+model_cfg = dict(
+    type='CDCv2',
+    name='rodnet-cdc-v2-win16-wobg',
+    max_dets=20,
+    peak_thres=0.3,
+    ols_thres=0.3,
+    mnet_cfg=(4, 32),
+    dcn=True,
+)
+
+confmap_cfg = dict(
+    confmap_sigmas={
+        'pedestrian': 15,
+        'cyclist': 20,
+        'car': 30,
+        # 'van': 40,
+        # 'truck': 50,
+    },
+    confmap_sigmas_interval={
+        'pedestrian': [5, 15],
+        'cyclist': [8, 20],
+        'car': [10, 30],
+        # 'van': [15, 40],
+        # 'truck': [20, 50],
+    },
+    confmap_length={
+        'pedestrian': 1,
+        'cyclist': 2,
+        'car': 3,
+        # 'van': 4,
+        # 'truck': 5,
+    }
+)
+
+train_cfg = dict(
+    n_epoch=100,
+    batch_size=4,
+    lr=0.00001,
+    lr_step=5,  # lr will decrease 10 times after lr_step epoches
+    win_size=16,
+    train_step=1,
+    train_stride=4,
+    log_step=100,
+    save_step=10000,
+)
+test_cfg = dict(
+    test_step=1,
+    test_stride=8,
+    rr_min=1.0,  # min radar range
+    rr_max=20.0,  # max radar range
+    ra_min=-60.0,  # min radar angle
+    ra_max=60.0,  # max radar angle
+)

configs/config_rodnet_v2_hg1_win16.py

+dataset_cfg = dict(
+    dataset_name='ROD2021',
+    base_root="/mnt/disk2/CRUW/ROD2021",
+    data_root="/mnt/disk2/CRUW/ROD2021/sequences",
+    anno_root="/mnt/disk2/CRUW/ROD2021/annotations",
+    anno_ext='.txt',
+    train=dict(
+        subdir='train',
+        # seqs=[],  # can choose from the subdir folder
+    ),
+    valid=dict(
+        subdir='valid',
+        seqs=[],
+    ),
+    test=dict(
+        subdir='test',
+        # seqs=[],  # can choose from the subdir folder
+    ),
+    demo=dict(
+        subdir='demo',
+        seqs=[],
+    ),
+)
+
+model_cfg = dict(
+    type='HGv2',
+    name='rodnet-hg1-v2-win16-wobg',
+    max_dets=20,
+    peak_thres=0.4,
+    ols_thres=0.3,
+    stacked_num=1,
+    mnet_cfg=(4, 32),
+    dcn=True,
+)
+
+confmap_cfg = dict(
+    confmap_sigmas={
+        'pedestrian': 15,
+        'cyclist': 20,
+        'car': 30,
+        # 'van': 40,
+        # 'truck': 50,
+    },
+    confmap_sigmas_interval={
+        'pedestrian': [5, 15],
+        'cyclist': [8, 20],
+        'car': [10, 30],
+        # 'van': [15, 40],
+        # 'truck': [20, 50],
+    },
+    confmap_length={
+        'pedestrian': 1,
+        'cyclist': 2,
+        'car': 3,
+        # 'van': 4,
+        # 'truck': 5,
+    }
+)
+
+train_cfg = dict(
+    n_epoch=50,
+    batch_size=4,
+    lr=0.00001,
+    lr_step=5,  # lr will decrease 10 times after lr_step epoches
+    win_size=16,
+    train_step=1,
+    train_stride=4,
+    log_step=100,
+    save_step=10000,
+)
+test_cfg = dict(
+    test_step=1,
+    test_stride=8,
+    rr_min=1.0,  # min radar range
+    rr_max=20.0,  # max radar range
+    ra_min=-60.0,  # min radar angle
+    ra_max=60.0,  # max radar angle
+)

configs/config_rodnet_v2_hg1wi_win16.py

+dataset_cfg = dict(
+    dataset_name='ROD2021',
+    base_root="/mnt/disk1/CRUW/ROD2021",
+    data_root="/mnt/disk1/CRUW/ROD2021/sequences",
+    anno_root="/mnt/disk1/CRUW/ROD2021/annotations",
+    anno_ext='.txt',
+    train=dict(
+        subdir='train',
+        # seqs=[],  # can choose from the subdir folder
+    ),
+    valid=dict(
+        subdir='valid',
+        # seqs=[],
+    ),
+    test=dict(
+        subdir='test',
+        # seqs=[],  # can choose from the subdir folder
+    ),
+    demo=dict(
+        subdir='demo',
+        seqs=[],
+    ),
+)
+
+model_cfg = dict(
+    type='HGwIv2',
+    name='rodnet-hg1wi-v2-win16-wobg',
+    max_dets=20,
+    peak_thres=0.3,
+    ols_thres=0.3,
+    stacked_num=1,
+    mnet_cfg=(4, 32),
+    dcn=True,
+)
+
+confmap_cfg = dict(
+    confmap_sigmas={
+        'pedestrian': 15,
+        'cyclist': 20,
+        'car': 30,
+        # 'van': 40,
+        # 'truck': 50,
+    },
+    confmap_sigmas_interval={
+        'pedestrian': [5, 15],
+        'cyclist': [8, 20],
+        'car': [10, 30],
+        # 'van': [15, 40],
+        # 'truck': [20, 50],
+    },
+    confmap_length={
+        'pedestrian': 1,
+        'cyclist': 2,
+        'car': 3,
+        # 'van': 4,
+        # 'truck': 5,
+    }
+)
+
+train_cfg = dict(
+    n_epoch=100,
+    batch_size=2,
+    lr=0.00001,
+    lr_step=5,  # lr will decrease 10 times after lr_step epoches
+    win_size=16,
+    train_step=1,
+    train_stride=4,
+    log_step=100,
+    save_step=10000,
+)
+test_cfg = dict(
+    test_step=1,
+    test_stride=8,
+    rr_min=1.0,  # min radar range
+    rr_max=20.0,  # max radar range
+    ra_min=-60.0,  # min radar angle
+    ra_max=60.0,  # max radar angle
+)

rodnet/datasets/CRDataLoader.py

@@ -5,14 +5,14 @@ import torch
 from multiprocessing import Array, Process
 
 
-class CRDataLoader():
+class CRDataLoader:
     def __init__(self, dataset, shuffle=False, num_parallel_batch=2, noise_channel=False):
         # parameters settings
         self.dataset = dataset
         self.config_dict = self.dataset.config_dict
-        self.n_class = self.config_dict['class_cfg']['n_class']
+        self.n_class = dataset.dataset.object_cfg.n_class
         self.batch_size = self.config_dict['train_cfg']['batch_size']
-        self.radar_configs = self.config_dict['dataset_cfg']['radar_cfg']
+        self.radar_configs = self.dataset.dataset.sensor_cfg.radar_cfg
         self.model_configs = self.config_dict['model_cfg']
         self.ramap_rsize = self.radar_configs['ramap_rsize']
         self.ramap_asize = self.radar_configs['ramap_asize']

rodnet/datasets/CRDataset.py

@@ -45,10 +45,7 @@ class CRDataset(data.Dataset):
         if 'mnet_cfg' in self.config_dict['model_cfg']:
             in_chirps, out_channels = self.config_dict['model_cfg']['mnet_cfg']
             self.n_chirps = in_chirps
-            n_radar_chirps = self.config_dict['dataset_cfg']['radar_cfg']['n_chirps']
-            self.chirp_ids = []
-            for c in range(in_chirps):
-                self.chirp_ids.append(int(n_radar_chirps / in_chirps * c))
+        self.chirp_ids = self.dataset.sensor_cfg.radar_cfg['chirp_ids']
 
         # dataset initialization
         self.image_paths = []
@@ -105,7 +102,7 @@ class CRDataset(data.Dataset):
         )
 
         if self.is_random_chirp:
-            chirp_id = random.randint(0, self.dataset.sensor_cfg.radar_cfg['n_chirps'] - 1)
+            chirp_id = random.randint(0, len(self.chirp_ids) - 1)
         else:
             chirp_id = 0
 
@@ -144,12 +141,23 @@ class CRDataset(data.Dataset):
                 else:
                     raise TypeError
             elif radar_configs['data_type'] == 'ROD2021':
+                if isinstance(chirp_id, int):
                     radar_npy_win = np.zeros((self.win_size, ramap_rsize, ramap_asize, 2), dtype=np.float32)
-                chirp_id = 0  # only use chirp 0 for training
                     for idx, frameid in enumerate(
                             range(data_id, data_id + self.win_size * self.step, self.step)):
                         radar_npy_win[idx, :, :, :] = np.load(radar_paths[frameid][chirp_id])
                         data_dict['image_paths'].append(image_paths[frameid])
+                elif isinstance(chirp_id, list):
+                    radar_npy_win = np.zeros((self.win_size, self.n_chirps, ramap_rsize, ramap_asize, 2),
+                                             dtype=np.float32)
+                    for idx, frameid in enumerate(
+                            range(data_id, data_id + self.win_size * self.step, self.step)):
+                        for cid, c in enumerate(chirp_id):
+                            npy_path = radar_paths[frameid][cid]
+                            radar_npy_win[idx, cid, :, :, :] = np.load(npy_path)
+                        data_dict['image_paths'].append(image_paths[frameid])
+                else:
+                    raise TypeError
             else:
                 raise NotImplementedError
         except:
@@ -161,21 +169,6 @@ class CRDataset(data.Dataset):
             with open(os.path.join('./tmp', log_name), 'w') as f_log:
                 f_log.write('npy path: ' + radar_paths[frameid][chirp_id] + \
                             '\nframe indices: %d:%d:%d' % (data_id, data_id + self.win_size * self.step, self.step))
-            # radar_npy_win = np.transpose(radar_npy_win, (3, 0, 1, 2))
-            #
-            # data_dict['radar_data'] = radar_npy_win
-            #
-            # if len(self.confmaps) != 0:
-            #     confmap_gt = this_seq_confmap[data_id:data_id + self.win_size * self.step:self.step]
-            #     confmap_gt = np.transpose(confmap_gt, (1, 0, 2, 3))
-            #     obj_info = this_seq_obj_info[data_id:data_id + self.win_size * self.step:self.step]
-            #
-            #     data_dict['anno'] = dict(
-            #         obj_infos=obj_info,
-            #         confmaps=confmap_gt,
-            #     )
-            # else:
-            #     data_dict['anno'] = None
             return data_dict
 
         # Dataloader for MNet

rodnet/datasets/CRDatasetSM.py

@@ -8,7 +8,7 @@ from tqdm import tqdm
 from torch.utils import data
 
 from .collate_functions import _cr_collate_npy
-from .loaders import list_pkl_filenames
+from .loaders import list_pkl_filenames, list_pkl_filenames_from_prepared
 
 
 class CRDatasetSM(data.Dataset):
@@ -24,11 +24,12 @@ class CRDatasetSM(data.Dataset):
     :param is_random: random load or not
     """
 
-    def __init__(self, data_root, config_dict, split, is_random=True, subset=None, noise_channel=False):
+    def __init__(self, data_dir, dataset, config_dict, split, is_random=True, subset=None, noise_channel=False):
         # parameters settings
-        self.data_root = data_root
+        self.data_dir = data_dir
+        self.dataset = dataset
         self.config_dict = config_dict
-        self.n_class = config_dict['class_cfg']['n_class']
+        self.n_class = dataset.object_cfg.n_class
         self.win_size = config_dict['train_cfg']['win_size']
         self.split = split
         if split == 'train' or split == 'valid':
@@ -45,10 +46,7 @@ class CRDatasetSM(data.Dataset):
         if 'mnet_cfg' in self.config_dict['model_cfg']:
             in_chirps, out_channels = self.config_dict['model_cfg']['mnet_cfg']
             self.n_chirps = in_chirps
-            n_radar_chirps = self.config_dict['dataset_cfg']['radar_cfg']['n_chirps']
-            self.chirp_ids = []
-            for c in range(in_chirps):
-                self.chirp_ids.append(int(n_radar_chirps / in_chirps * c))
+        self.chirp_ids = self.dataset.sensor_cfg.radar_cfg['chirp_ids']
 
         # dataset initialization
         self.image_paths = []
@@ -60,12 +58,13 @@ class CRDatasetSM(data.Dataset):
         if subset is not None:
             self.data_files = [subset + '.pkl']
         else:
-            self.data_files = list_pkl_filenames(config_dict['dataset_cfg'], split)
+            # self.data_files = list_pkl_filenames(config_dict['dataset_cfg'], split)
+            self.data_files = list_pkl_filenames_from_prepared(data_dir, split)
         self.seq_names = [name.split('.')[0] for name in self.data_files]
         self.n_seq = len(self.seq_names)
 
         for seq_id, data_file in enumerate(tqdm(self.data_files)):
-            data_file_path = os.path.join(data_root, split, data_file)
+            data_file_path = os.path.join(data_dir, split, data_file)
             data_details = pickle.load(open(data_file_path, 'rb'))
             if split == 'train' or split == 'valid':
                 assert data_details['anno'] is not None
@@ -78,7 +77,7 @@ class CRDatasetSM(data.Dataset):
             for data_id in range(n_data_in_seq):
                 self.index_mapping.append([seq_id, data_id * self.stride])
             if data_details['anno'] is not None:
-                self.obj_infos.append(data_details['anno']['obj_infos'])
+                self.obj_infos.append(data_details['anno']['metadata'])
 
     def __len__(self):
         """Total number of data/label pairs"""
@@ -89,10 +88,10 @@ class CRDatasetSM(data.Dataset):
         seq_id, data_id = self.index_mapping[index]
         image_paths = self.image_paths[seq_id]
         radar_paths = self.radar_paths[seq_id]
-        this_data_file = os.path.join(self.data_root, self.split, self.data_files[seq_id])
+        this_data_file = os.path.join(self.data_dir, self.split, self.data_files[seq_id])
         this_data_details = pickle.load(open(this_data_file, 'rb'))
         if this_data_details['anno'] is not None:
-            this_seq_obj_info = this_data_details['anno']['obj_infos']
+            this_seq_obj_info = this_data_details['anno']['metadata']
             this_seq_confmap = this_data_details['anno']['confmaps']
 
         data_dict = dict(
@@ -101,7 +100,7 @@ class CRDatasetSM(data.Dataset):
         )
 
         if self.is_random:
-            chirp_id = random.randint(0, self.config_dict['dataset_cfg']['radar_cfg']['n_chirps'] - 1)
+            chirp_id = random.randint(0, len(self.chirp_ids) - 1)
         else:
             chirp_id = 0
 
@@ -109,7 +108,7 @@ class CRDatasetSM(data.Dataset):
         if 'mnet_cfg' in self.config_dict['model_cfg']:
             chirp_id = self.chirp_ids
 
-        radar_configs = self.config_dict['dataset_cfg']['radar_cfg']
+        radar_configs = self.dataset.sensor_cfg.radar_cfg
         ramap_rsize = radar_configs['ramap_rsize']
         ramap_asize = radar_configs['ramap_asize']
 
@@ -139,6 +138,24 @@ class CRDatasetSM(data.Dataset):
                         data_dict['image_paths'].append(image_paths[frameid])
                 else:
                     raise TypeError
+            elif radar_configs['data_type'] == 'ROD2021':
+                if isinstance(chirp_id, int):
+                    radar_npy_win = np.zeros((self.win_size, ramap_rsize, ramap_asize, 2), dtype=np.float32)
+                    for idx, frameid in enumerate(
+                            range(data_id, data_id + self.win_size * self.step, self.step)):
+                        radar_npy_win[idx, :, :, :] = np.load(radar_paths[frameid][chirp_id])
+                        data_dict['image_paths'].append(image_paths[frameid])
+                elif isinstance(chirp_id, list):
+                    radar_npy_win = np.zeros((self.win_size, self.n_chirps, ramap_rsize, ramap_asize, 2),
+                                             dtype=np.float32)
+                    for idx, frameid in enumerate(
+                            range(data_id, data_id + self.win_size * self.step, self.step)):
+                        for cid, c in enumerate(chirp_id):
+                            npy_path = radar_paths[frameid][cid]
+                            radar_npy_win[idx, cid, :, :, :] = np.load(npy_path)
+                        data_dict['image_paths'].append(image_paths[frameid])
+                else:
+                    raise TypeError
             else:
                 raise ValueError
 
@@ -151,24 +168,6 @@ class CRDatasetSM(data.Dataset):
             with open(os.path.join('./tmp', log_name), 'w') as f_log:
                 f_log.write('npy path: ' + radar_paths[frameid][chirp_id] + \
                             '\nframe indices: %d:%d:%d' % (data_id, data_id + self.win_size * self.step, self.step))
-            # if 'mnet_cfg' in self.config_dict['model_cfg']:
-            #     radar_npy_win = np.transpose(radar_npy_win, (4, 0, 1, 2, 3))
-            # else:
-            #     radar_npy_win = np.transpose(radar_npy_win, (3, 0, 1, 2))
-            #
-            # data_dict['radar_data'] = radar_npy_win
-            #
-            # if len(self.confmaps) != 0:
-            #     confmap_gt = this_seq_confmap[data_id:data_id + self.win_size * self.step:self.step]
-            #     confmap_gt = np.transpose(confmap_gt, (1, 0, 2, 3))
-            #     obj_info = this_seq_obj_info[data_id:data_id + self.win_size * self.step:self.step]
-            #
-            #     data_dict['anno'] = dict(
-            #         obj_infos=obj_info,
-            #         confmaps=confmap_gt,
-            #     )
-            # else:
-            #     data_dict['anno'] = None
             return data_dict
 
         # Dataloader for MNet

rodnet/models/__init__.py

 from .rodnet_cdc import RODNetCDC
 from .rodnet_hg import RODNetHG
 from .rodnet_hgwi import RODNetHGwI
+
+from .rodnet_cdc_v2 import RODNetCDCDCN
+from .rodnet_hg_v2 import RODNetHGDCN
+from .rodnet_hgwi_v2 import RODNetHGwIDCN

rodnet/models/backbones/cdc.py

 import torch.nn as nn
 
 
+class RadarVanilla(nn.Module):
+
+    def __init__(self, in_channels, n_class, use_mse_loss=False):
+        super(RadarVanilla, self).__init__()
+        self.encoder = RODEncode(in_channels=in_channels)
+        self.decoder = RODDecode(n_class=n_class)
+        self.sigmoid = nn.Sigmoid()
+        self.use_mse_loss = use_mse_loss
+
+    def forward(self, x):
+        x = self.encoder(x)
+        x = self.decoder(x)
+        if not self.use_mse_loss:
+            x = self.sigmoid(x)
+        return x
+
+
 class RODEncode(nn.Module):
 
     def __init__(self, in_channels=2):
@@ -32,7 +49,6 @@ class RODEncode(nn.Module):
         x = self.relu(self.bn2b(self.conv2b(x)))  # (B, 128, W/2, 64, 64) -> (B, 128, W/4, 32, 32)
         x = self.relu(self.bn3a(self.conv3a(x)))  # (B, 128, W/4, 32, 32) -> (B, 256, W/4, 32, 32)
         x = self.relu(self.bn3b(self.conv3b(x)))  # (B, 256, W/4, 32, 32) -> (B, 256, W/4, 16, 16)
-
         return x
 
 

rodnet/models/backbones/hg.py

@@ -3,11 +3,16 @@ import torch.nn as nn
 
 class RadarStackedHourglass(nn.Module):
 
-    def __init__(self, n_class, stacked_num=1, in_channels=2):
+    def __init__(self, in_channels, n_class, stacked_num=1, conv_op=None, use_mse_loss=False):
         super(RadarStackedHourglass, self).__init__()
         self.stacked_num = stacked_num
+        if conv_op is None:
             self.conv1a = nn.Conv3d(in_channels=in_channels, out_channels=32,
                                     kernel_size=(9, 5, 5), stride=(1, 1, 1), padding=(4, 2, 2))
+        else:
+            self.conv1a = conv_op(in_channels=in_channels, out_channels=32,
+                                  kernel_size=(5, 3, 3), stride=(1, 1, 1), padding=(2, 1, 1))
+
         self.conv1b = nn.Conv3d(in_channels=32, out_channels=64,
                                 kernel_size=(9, 5, 5), stride=(1, 1, 1), padding=(4, 2, 2))
 
@@ -25,6 +30,7 @@ class RadarStackedHourglass(nn.Module):
         self.bn1a = nn.BatchNorm3d(num_features=32)
         self.bn1b = nn.BatchNorm3d(num_features=64)
         self.sigmoid = nn.Sigmoid()
+        self.use_mse_loss = use_mse_loss
 
     def forward(self, x):
         x = self.relu(self.bn1a(self.conv1a(x)))
@@ -35,11 +41,12 @@ class RadarStackedHourglass(nn.Module):
             x, x1, x2, x3 = self.hourglass[i][0](x)
             x = self.hourglass[i][1](x, x1, x2, x3)
             confmap = self.hourglass[i][2](x)
-            out.append(self.sigmoid(confmap))
+            if not self.use_mse_loss:
+                confmap = self.sigmoid(confmap)
+            out.append(confmap)
             if i < self.stacked_num - 1:
                 confmap_ = self.hourglass[i][3](confmap)
                 x = x + confmap_
-
         return out
 
 

rodnet/models/backbones/hgwi.py

@@ -4,11 +4,16 @@ import torch.nn as nn
 
 class RadarStackedHourglass(nn.Module):
 
-    def __init__(self, n_class, stacked_num=1):
+    def __init__(self, in_channels, n_class, stacked_num=1, conv_op=None, use_mse_loss=False):
         super(RadarStackedHourglass, self).__init__()
         self.stacked_num = stacked_num
-        self.conv1a = nn.Conv3d(in_channels=2, out_channels=32,
+        if conv_op is None:
+            self.conv1a = nn.Conv3d(in_channels=in_channels, out_channels=32,
                                     kernel_size=(9, 5, 5), stride=(1, 1, 1), padding=(4, 2, 2))
+        else:
+            self.conv1a = conv_op(in_channels=in_channels, out_channels=32,
+                                  kernel_size=(5, 3, 3), stride=(1, 1, 1), padding=(2, 1, 1))
+
         self.conv1b = nn.Conv3d(in_channels=32, out_channels=64,
                                 kernel_size=(9, 5, 5), stride=(1, 1, 1), padding=(4, 2, 2))
         self.conv1c = nn.Conv3d(in_channels=64, out_channels=160,
@@ -29,6 +34,7 @@ class RadarStackedHourglass(nn.Module):
         self.bn1b = nn.BatchNorm3d(num_features=64)
         self.bn1c = nn.BatchNorm3d(num_features=160)
         self.sigmoid = nn.Sigmoid()
+        self.use_mse_loss = use_mse_loss
 
     def forward(self, x):
         x = self.relu(self.bn1a(self.conv1a(x)))
@@ -40,11 +46,12 @@ class RadarStackedHourglass(nn.Module):
             x, x1, x2, x3 = self.hourglass[i][0](x)
             x = self.hourglass[i][1](x, x1, x2, x3)
             confmap = self.hourglass[i][2](x)
-            out.append(self.sigmoid(confmap))
+            if not self.use_mse_loss:
+                confmap = self.sigmoid(confmap)
+            out.append(confmap)
             if i < self.stacked_num - 1:
                 confmap_ = self.hourglass[i][3](confmap)
                 x = x + confmap_
-
         return out
 
 

rodnet/models/modules/mnet.py

+import math
+import torch
+import torch.nn as nn
+
+
+class MNet(nn.Module):
+    def __init__(self, in_chirps, out_channels, conv_op=None):
+        super(MNet, self).__init__()
+        self.in_chirps = in_chirps
+        self.out_channels = out_channels
+        if conv_op is None:
+            conv_op = nn.Conv3d
+        self.conv_op = conv_op
+
+        self.t_conv3d = conv_op(in_channels=2, out_channels=out_channels, kernel_size=(3, 1, 1), stride=(2, 1, 1),
+                                padding=(1, 0, 0))
+        t_conv_out = math.floor((in_chirps + 2 * 1 - (3 - 1) - 1) / 2 + 1)
+        self.t_maxpool = nn.MaxPool3d(kernel_size=(t_conv_out, 1, 1))
+
+    def forward(self, x):
+        batch_size, n_channels, win_size, in_chirps, w, h = x.shape
+        x_out = torch.zeros((batch_size, self.out_channels, win_size, w, h)).cuda()
+        for win in range(win_size):
+            x_win = self.t_conv3d(x[:, :, win, :, :, :])
+            x_win = self.t_maxpool(x_win)
+            x_win = x_win.view(batch_size, self.out_channels, w, h)
+            x_out[:, :, win, ] = x_win
+        return x_out
+
+
+if __name__ == '__main__':
+    batch_size = 4
+    in_channels = 2
+    win_size = 32
+    in_chirps = 4
+    w = 128
+    h = 128
+    out_channels = 32
+    mnet = MNet(in_chirps=in_chirps, out_channels=out_channels)
+    input = torch.randn(batch_size, in_channels, win_size, in_chirps, w, h)
+    output = mnet(input)

rodnet/models/rodnet_cdc.py

 import torch.nn as nn
 
-from .backbones.cdc import RODEncode, RODDecode
+from .backbones.cdc import RadarVanilla
 
 
 class RODNetCDC(nn.Module):
-    def __init__(self, n_class):
+    def __init__(self, in_channels, n_class):
         super(RODNetCDC, self).__init__()
-
-        self.c3d_encode = RODEncode()
-        self.c3d_decode = RODDecode(n_class)
+        self.cdc = RadarVanilla(in_channels, n_class, use_mse_loss=False)
 
     def forward(self, x):
-        x = self.c3d_encode(x)
-        dets = self.c3d_decode(x)
-        return dets
+        x = self.cdc(x)
+        return x

rodnet/models/rodnet_cdc_v2.py

+import torch.nn as nn
+
+from .backbones.cdc import RadarVanilla
+from .modules.mnet import MNet
+from ..ops.dcn import DeformConvPack3D
+
+
+class RODNetCDCDCN(nn.Module):
+    def __init__(self, in_channels, n_class, mnet_cfg=None, dcn=True):
+        super(RODNetCDCDCN, self).__init__()
+        self.dcn = dcn
+        if dcn:
+            self.conv_op = DeformConvPack3D
+        else:
+            self.conv_op = nn.Conv3d
+        if mnet_cfg is not None:
+            in_chirps_mnet, out_channels_mnet = mnet_cfg
+            assert in_channels == in_chirps_mnet
+            self.mnet = MNet(in_chirps_mnet, out_channels_mnet, conv_op=self.conv_op)
+            self.with_mnet = True
+            self.cdc = RadarVanilla(out_channels_mnet, n_class, use_mse_loss=False)
+        else:
+            self.with_mnet = False
+            self.cdc = RadarVanilla(in_channels, n_class, use_mse_loss=False)
+
+    def forward(self, x):
+        if self.with_mnet:
+            x = self.mnet(x)
+        x = self.cdc(x)
+        return x

rodnet/models/rodnet_hg.py

@@ -4,9 +4,9 @@ from .backbones.hg import RadarStackedHourglass
 
 
 class RODNetHG(nn.Module):
-    def __init__(self, n_class, stacked_num=2):
+    def __init__(self, in_channels, n_class, stacked_num=2):
         super(RODNetHG, self).__init__()
-        self.stacked_hourglass = RadarStackedHourglass(n_class, stacked_num=stacked_num)
+        self.stacked_hourglass = RadarStackedHourglass(in_channels, n_class, stacked_num=stacked_num)
 
     def forward(self, x):
         out = self.stacked_hourglass(x)

rodnet/models/rodnet_hg_v2.py

+import torch
+import torch.nn as nn
+
+from .backbones.hg import RadarStackedHourglass
+from .modules.mnet import MNet
+from ..ops.dcn import DeformConvPack3D
+
+
+class RODNetHGDCN(nn.Module):
+    def __init__(self, in_channels, n_class, stacked_num=2, mnet_cfg=None, dcn=True):
+        super(RODNetHGDCN, self).__init__()
+        self.dcn = dcn
+        if dcn:
+            self.conv_op = DeformConvPack3D
+        else:
+            self.conv_op = nn.Conv3d
+        if mnet_cfg is not None:
+            in_chirps_mnet, out_channels_mnet = mnet_cfg
+            assert in_channels == in_chirps_mnet
+            self.mnet = MNet(in_chirps_mnet, out_channels_mnet, conv_op=self.conv_op)
+            self.with_mnet = True
+            self.stacked_hourglass = RadarStackedHourglass(out_channels_mnet, n_class, stacked_num=stacked_num,
+                                                           conv_op=self.conv_op)
+        else:
+            self.with_mnet = False
+            self.stacked_hourglass = RadarStackedHourglass(in_channels, n_class, stacked_num=stacked_num,
+                                                           conv_op=self.conv_op)
+
+    def forward(self, x):
+        if self.with_mnet:
+            x = self.mnet(x)
+        out = self.stacked_hourglass(x)
+        return out
+
+
+if __name__ == '__main__':
+    torch.cuda.set_device(torch.device('cuda:0'))
+    batch_size = 1
+    in_channels = 2
+    win_size = 4
+    in_chirps = 6
+    w = 64
+    h = 64
+    out_channels = 8
+    model = RODNetHGDCN(n_class=3, stacked_num=1, mnet_cfg=(in_chirps, out_channels)).cuda()
+    for iter in range(10):
+        input = torch.randn(batch_size, in_channels, win_size, in_chirps, w, h).cuda()
+        output = model(input)
+        print("forward done")
+        output_gt = torch.randn(batch_size, 3, win_size, w, h).cuda()
+        criterion = nn.BCELoss()
+        loss = criterion(output[0], output_gt)
+        loss.backward()

rodnet/models/rodnet_hgwi.py

@@ -5,9 +5,9 @@ from .backbones.hgwi import RadarStackedHourglass
 
 
 class RODNetHGwI(nn.Module):
-    def __init__(self, n_class, stacked_num=1):
+    def __init__(self, in_channels, n_class, stacked_num=1):
         super(RODNetHGwI, self).__init__()
-        self.stacked_hourglass = RadarStackedHourglass(n_class, stacked_num=stacked_num)
+        self.stacked_hourglass = RadarStackedHourglass(in_channels, n_class, stacked_num=stacked_num)
 
     def forward(self, x):
         out = self.stacked_hourglass(x)

rodnet/models/rodnet_hgwi_v2.py

+import torch
+import torch.nn as nn
+
+from .backbones.hgwi import RadarStackedHourglass
+from .modules.mnet import MNet
+from ..ops.dcn import DeformConvPack3D
+
+
+class RODNetHGwIDCN(nn.Module):
+    def __init__(self, in_channels, n_class, stacked_num=1, mnet_cfg=None, dcn=True):
+        super(RODNetHGwIDCN, self).__init__()
+        self.dcn = dcn
+        if dcn:
+            self.conv_op = DeformConvPack3D
+        else:
+            self.conv_op = nn.Conv3d
+        if mnet_cfg is not None:
+            in_chirps_mnet, out_channels_mnet = mnet_cfg
+            self.mnet = MNet(in_chirps_mnet, out_channels_mnet, conv_op=self.conv_op)
+            self.with_mnet = True
+            self.stacked_hourglass = RadarStackedHourglass(out_channels_mnet, n_class, stacked_num=stacked_num,
+                                                           conv_op=self.conv_op)
+        else:
+            self.with_mnet = False
+            self.stacked_hourglass = RadarStackedHourglass(in_channels, n_class, stacked_num=stacked_num,
+                                                           conv_op=self.conv_op)
+
+    def forward(self, x):
+        if self.with_mnet:
+            x = self.mnet(x)
+        out = self.stacked_hourglass(x)
+        return out
+
+
+if __name__ == '__main__':
+    testModel = RODNetHGwIDCN().cuda()
+    x = torch.zeros((1, 2, 16, 128, 128)).cuda()
+    testModel(x)