Upload 30 files

Dockerfile

@@ -0,0 +1,14 @@
+# read the doc: https://huggingface.co/docs/hub/spaces-sdks-docker
+# you will also find guides on how best to write your Dockerfile
+
+FROM python:3.9
+
+WORKDIR /code
+
+COPY ./requirements.txt /code/requirements.txt
+
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+
+COPY . .
+
+CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]

code/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2020 CVLab@StonyBrook
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

code/README.md

@@ -0,0 +1,83 @@
+# DM-Count
+
+Official Pytorch implementation of the paper Distribution Matching for Crowd Counting (NeurIPS, spotlight).
+
+[Arxiv](https://arxiv.org/pdf/2009.13077.pdf) | [NeurIPS Processings](https://proceedings.neurips.cc/paper/2020/hash/118bd558033a1016fcc82560c65cca5f-Abstract.html)
+
+We propose to use Distribution Matching for crowd COUNTing (DM-Count). In DM-Count, we use Optimal Transport (OT) to measure the similarity between the normalized predicted density map and the normalized ground truth density map. To stabilize OT computation, we include a Total Variation loss in our model. We show that the generalization error bound of DM-Count is tighter than that of the Gaussian smoothed methods. Empirically, our method outperforms the state-of-the-art methods by a large margin on four challenging crowd counting datasets: UCF-QNRF, NWPU, ShanghaiTech, and UCF-CC50.
+
+
+## Prerequisites
+
+Python 3.x
+
+Pytorch >= 1.2
+
+For other libraries, check requirements.txt.
+
+## Getting Started
+1. Dataset download
+
++ QNRF can be downloaded [here](https://www.crcv.ucf.edu/data/ucf-qnrf/)
+
++ NWPU can be downloaded [here](https://www.crowdbenchmark.com/nwpucrowd.html)
+
++ Shanghai Tech Part A and Part B can be downloaded [here](https://www.kaggle.com/tthien/shanghaitech)
+
+2. Data preprocess
+
+Due to large sizes of images in QNRF and NWPU datasets, we preprocess these two datasets.
+
+```
+python preprocess_dataset.py --dataset <dataset name: qnrf or nwpu> --input-dataset-path <original data directory> --output-dataset-path <processed data directory> 
+```
+    
+3. Training
+
+```
+python train.py --dataset <dataset name: qnrf, sha, shb or nwpu> --data-dir <path to dataset> --device <gpu device id>
+```
+
+4. Test
+
+```
+python test.py --model-path <path of the model to be evaluated> --data-path <directory for the dataset> --dataset <dataset name: qnrf, sha, shb or nwpu>
+```
+
+## Pretrained models
+
+Pretrained models on UCF-QNRF, NWPU, Shanghaitech part A and B can be found [Google Drive](https://drive.google.com/drive/folders/10U7F4iW_aPICM5-qJq21SXLLkzlum9tX?usp=sharing). You could download them and put them in in pretrained_models folder.
+
+
+## Other resources
+
++ Web Demo
+
+A web interface to can be found [here](https://gradio.app/g/dm-count). 
+
+![demo](https://i.ibb.co/m65HpCJ/dm-count.gif)
+
+Feel free to upload a image and try out the demo on a web browser. It is developed by [Ali Abdalla](twitter.com/si3luwa) from [Gradio](https://github.com/gradio-app/gradio). Gradio is an open source library, which helps to create interfaces to make models more accessible. Thanks Ali and Gradio! 
+
+To launch a Gradio interface, run 
+
+```
+ python demo.py
+```
+
++ Kaggle Notebook
+
+A [Kaggle Notebook](https://www.kaggle.com/selmanzleyen/dmcount-shb) is developed by [Selman Ozleyen](https://github.com/SelmanOzleyen/DM-Count). Thanks Selman!
+
+
+## References
+If you find this work or code useful, please cite:
+
+```
+@inproceedings{wang2020DMCount,
+  title={Distribution Matching for Crowd Counting},
+  author={Boyu Wang and Huidong Liu and Dimitris Samaras and Minh Hoai},
+  booktitle={Advances in Neural Information Processing Systems},
+  year={2020},
+}
+```

code/datasets/crowd.py

@@ -0,0 +1,234 @@
+from PIL import Image
+import torch.utils.data as data
+import os
+from glob import glob
+import torch
+import torchvision.transforms.functional as F
+from torchvision import transforms
+import random
+import numpy as np
+import scipy.io as sio
+
+
+def random_crop(im_h, im_w, crop_h, crop_w):
+    res_h = im_h - crop_h
+    res_w = im_w - crop_w
+    i = random.randint(0, res_h)
+    j = random.randint(0, res_w)
+    return i, j, crop_h, crop_w
+
+
+def gen_discrete_map(im_height, im_width, points):
+    """
+        func: generate the discrete map.
+        points: [num_gt, 2], for each row: [width, height]
+        """
+    discrete_map = np.zeros([im_height, im_width], dtype=np.float32)
+    h, w = discrete_map.shape[:2]
+    num_gt = points.shape[0]
+    if num_gt == 0:
+        return discrete_map
+    
+    # fast create discrete map
+    points_np = np.array(points).round().astype(int)
+    p_h = np.minimum(points_np[:, 1], np.array([h-1]*num_gt).astype(int))
+    p_w = np.minimum(points_np[:, 0], np.array([w-1]*num_gt).astype(int))
+    p_index = torch.from_numpy(p_h* im_width + p_w)
+    discrete_map = torch.zeros(im_width * im_height).scatter_add_(0, index=p_index, src=torch.ones(im_width*im_height)).view(im_height, im_width).numpy()
+
+    ''' slow method
+    for p in points:
+        p = np.round(p).astype(int)
+        p[0], p[1] = min(h - 1, p[1]), min(w - 1, p[0])
+        discrete_map[p[0], p[1]] += 1
+    '''
+    assert np.sum(discrete_map) == num_gt
+    return discrete_map
+
+
+class Base(data.Dataset):
+    def __init__(self, root_path, crop_size, downsample_ratio=8):
+
+        self.root_path = root_path
+        self.c_size = crop_size
+        self.d_ratio = downsample_ratio
+        assert self.c_size % self.d_ratio == 0
+        self.dc_size = self.c_size // self.d_ratio
+        self.trans = transforms.Compose([
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+        ])
+
+    def __len__(self):
+        pass
+
+    def __getitem__(self, item):
+        pass
+
+    def train_transform(self, img, keypoints):
+        wd, ht = img.size
+        st_size = 1.0 * min(wd, ht)
+        assert st_size >= self.c_size
+        assert len(keypoints) >= 0
+        i, j, h, w = random_crop(ht, wd, self.c_size, self.c_size)
+        img = F.crop(img, i, j, h, w)
+        if len(keypoints) > 0:
+            keypoints = keypoints - [j, i]
+            idx_mask = (keypoints[:, 0] >= 0) * (keypoints[:, 0] <= w) * \
+                       (keypoints[:, 1] >= 0) * (keypoints[:, 1] <= h)
+            keypoints = keypoints[idx_mask]
+        else:
+            keypoints = np.empty([0, 2])
+
+        gt_discrete = gen_discrete_map(h, w, keypoints)
+        down_w = w // self.d_ratio
+        down_h = h // self.d_ratio
+        gt_discrete = gt_discrete.reshape([down_h, self.d_ratio, down_w, self.d_ratio]).sum(axis=(1, 3))
+        assert np.sum(gt_discrete) == len(keypoints)
+
+        if len(keypoints) > 0:
+            if random.random() > 0.5:
+                img = F.hflip(img)
+                gt_discrete = np.fliplr(gt_discrete)
+                keypoints[:, 0] = w - keypoints[:, 0]
+        else:
+            if random.random() > 0.5:
+                img = F.hflip(img)
+                gt_discrete = np.fliplr(gt_discrete)
+        gt_discrete = np.expand_dims(gt_discrete, 0)
+
+        return self.trans(img), torch.from_numpy(keypoints.copy()).float(), torch.from_numpy(
+            gt_discrete.copy()).float()
+
+
+class Crowd_qnrf(Base):
+    def __init__(self, root_path, crop_size,
+                 downsample_ratio=8,
+                 method='train'):
+        super().__init__(root_path, crop_size, downsample_ratio)
+        self.method = method
+        self.im_list = sorted(glob(os.path.join(self.root_path, '*.jpg')))
+        print('number of img: {}'.format(len(self.im_list)))
+        if method not in ['train', 'val']:
+            raise Exception("not implement")
+
+    def __len__(self):
+        return len(self.im_list)
+
+    def __getitem__(self, item):
+        img_path = self.im_list[item]
+        gd_path = img_path.replace('jpg', 'npy')
+        img = Image.open(img_path).convert('RGB')
+        if self.method == 'train':
+            keypoints = np.load(gd_path)
+            return self.train_transform(img, keypoints)
+        elif self.method == 'val':
+            keypoints = np.load(gd_path)
+            img = self.trans(img)
+            name = os.path.basename(img_path).split('.')[0]
+            return img, len(keypoints), name
+
+
+class Crowd_nwpu(Base):
+    def __init__(self, root_path, crop_size,
+                 downsample_ratio=8,
+                 method='train'):
+        super().__init__(root_path, crop_size, downsample_ratio)
+        self.method = method
+        self.im_list = sorted(glob(os.path.join(self.root_path, '*.jpg')))
+        print('number of img: {}'.format(len(self.im_list)))
+
+        if method not in ['train', 'val', 'test']:
+            raise Exception("not implement")
+
+    def __len__(self):
+        return len(self.im_list)
+
+    def __getitem__(self, item):
+        img_path = self.im_list[item]
+        gd_path = img_path.replace('jpg', 'npy')
+        img = Image.open(img_path).convert('RGB')
+        if self.method == 'train':
+            keypoints = np.load(gd_path)
+            return self.train_transform(img, keypoints)
+        elif self.method == 'val':
+            keypoints = np.load(gd_path)
+            img = self.trans(img)
+            name = os.path.basename(img_path).split('.')[0]
+            return img, len(keypoints), name
+        elif self.method == 'test':
+            img = self.trans(img)
+            name = os.path.basename(img_path).split('.')[0]
+            return img, name
+
+
+class Crowd_sh(Base):
+    def __init__(self, root_path, crop_size,
+                 downsample_ratio=8,
+                 method='train'):
+        super().__init__(root_path, crop_size, downsample_ratio)
+        self.method = method
+        if method not in ['train', 'val']:
+            raise Exception("not implement")
+
+        self.im_list = sorted(glob(os.path.join(self.root_path, 'images', '*.jpg')))
+        print('number of img: {}'.format(len(self.im_list)))
+
+    def __len__(self):
+        return len(self.im_list)
+
+    def __getitem__(self, item):
+        img_path = self.im_list[item]
+        name = os.path.basename(img_path).split('.')[0]
+        gd_path = os.path.join(self.root_path, 'ground-truth', 'GT_{}.mat'.format(name))
+        img = Image.open(img_path).convert('RGB')
+        keypoints = sio.loadmat(gd_path)['image_info'][0][0][0][0][0]
+
+        if self.method == 'train':
+            return self.train_transform(img, keypoints)
+        elif self.method == 'val':
+            img = self.trans(img)
+            return img, len(keypoints), name
+
+    def train_transform(self, img, keypoints):
+        wd, ht = img.size
+        st_size = 1.0 * min(wd, ht)
+        # resize the image to fit the crop size
+        if st_size < self.c_size:
+            rr = 1.0 * self.c_size / st_size
+            wd = round(wd * rr)
+            ht = round(ht * rr)
+            st_size = 1.0 * min(wd, ht)
+            img = img.resize((wd, ht), Image.BICUBIC)
+            keypoints = keypoints * rr
+        assert st_size >= self.c_size, print(wd, ht)
+        assert len(keypoints) >= 0
+        i, j, h, w = random_crop(ht, wd, self.c_size, self.c_size)
+        img = F.crop(img, i, j, h, w)
+        if len(keypoints) > 0:
+            keypoints = keypoints - [j, i]
+            idx_mask = (keypoints[:, 0] >= 0) * (keypoints[:, 0] <= w) * \
+                       (keypoints[:, 1] >= 0) * (keypoints[:, 1] <= h)
+            keypoints = keypoints[idx_mask]
+        else:
+            keypoints = np.empty([0, 2])
+
+        gt_discrete = gen_discrete_map(h, w, keypoints)
+        down_w = w // self.d_ratio
+        down_h = h // self.d_ratio
+        gt_discrete = gt_discrete.reshape([down_h, self.d_ratio, down_w, self.d_ratio]).sum(axis=(1, 3))
+        assert np.sum(gt_discrete) == len(keypoints)
+
+        if len(keypoints) > 0:
+            if random.random() > 0.5:
+                img = F.hflip(img)
+                gt_discrete = np.fliplr(gt_discrete)
+                keypoints[:, 0] = w - keypoints[:, 0] - 1
+        else:
+            if random.random() > 0.5:
+                img = F.hflip(img)
+                gt_discrete = np.fliplr(gt_discrete)
+        gt_discrete = np.expand_dims(gt_discrete, 0)
+
+        return self.trans(img), torch.from_numpy(keypoints.copy()).float(), torch.from_numpy(
+            gt_discrete.copy()).float()

code/demo.py

@@ -0,0 +1,54 @@
+import torch
+from models import vgg19
+import gdown
+from PIL import Image
+from torchvision import transforms
+import gradio as gr
+import cv2
+import numpy as np
+import scipy
+
+model_path = "pretrained_models/model_qnrf.pth"
+url = "https://drive.google.com/uc?id=1nnIHPaV9RGqK8JHL645zmRvkNrahD9ru"
+gdown.download(url, model_path, quiet=False)
+
+device = torch.device('cpu')  # device can be "cpu" or "gpu"
+
+model = vgg19()
+model.to(device)
+model.load_state_dict(torch.load(model_path, device))
+model.eval()
+
+
+def predict(inp):
+    inp = Image.fromarray(inp.astype('uint8'), 'RGB')
+    inp = transforms.ToTensor()(inp).unsqueeze(0)
+    inp = inp.to(device)
+    with torch.set_grad_enabled(False):
+        outputs, _ = model(inp)
+    count = torch.sum(outputs).item()
+    vis_img = outputs[0, 0].cpu().numpy()
+    # normalize density map values from 0 to 1, then map it to 0-255.
+    vis_img = (vis_img - vis_img.min()) / (vis_img.max() - vis_img.min() + 1e-5)
+    vis_img = (vis_img * 255).astype(np.uint8)
+    vis_img = cv2.applyColorMap(vis_img, cv2.COLORMAP_JET)
+    vis_img = cv2.cvtColor(vis_img, cv2.COLOR_BGR2RGB)
+    return vis_img, int(count)
+
+
+inputs = gr.Image(label="Image of Crowd")
+outputs = [
+    gr.Image(label="Predicted Density Map"),
+    gr.Label(label="Predicted Count")
+]
+
+# Assuming `title`, `desc`, and `examples` variables are defined elsewhere in your code.
+title = "Your App Title"
+desc = "Your App Description"
+
+gr.Interface(fn=predict, 
+             inputs=inputs, 
+             outputs=outputs, 
+             title=title, 
+             description=desc, 
+             allow_flagging="never").launch(share=True)

code/losses/__init__.py

@@ -0,0 +1 @@
+

code/losses/bregman_pytorch.py

@@ -0,0 +1,484 @@
+# -*- coding: utf-8 -*-
+"""
+Rewrite ot.bregman.sinkhorn in Python Optimal Transport (https://pythonot.github.io/_modules/ot/bregman.html#sinkhorn)
+using pytorch operations.
+Bregman projections for regularized OT (Sinkhorn distance).
+"""
+
+import torch
+
+M_EPS = 1e-16
+
+
+def sinkhorn(a, b, C, reg=1e-1, method='sinkhorn', maxIter=1000, tau=1e3,
+             stopThr=1e-9, verbose=False, log=True, warm_start=None, eval_freq=10, print_freq=200, **kwargs):
+    """
+    Solve the entropic regularization optimal transport
+    The input should be PyTorch tensors
+    The function solves the following optimization problem:
+
+    .. math::
+        \gamma = arg\min_\gamma <\gamma,C>_F + reg\cdot\Omega(\gamma)
+        s.t. \gamma 1 = a
+             \gamma^T 1= b
+             \gamma\geq 0
+    where :
+    - C is the (ns,nt) metric cost matrix
+    - :math:`\Omega` is the entropic regularization term :math:`\Omega(\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})`
+    - a and b are target and source measures (sum to 1)
+    The algorithm used for solving the problem is the Sinkhorn-Knopp matrix scaling algorithm as proposed in [1].
+
+    Parameters
+    ----------
+    a : torch.tensor (na,)
+        samples measure in the target domain
+    b : torch.tensor (nb,)
+        samples in the source domain
+    C : torch.tensor (na,nb)
+        loss matrix
+    reg : float
+        Regularization term > 0
+    method : str
+        method used for the solver either 'sinkhorn', 'greenkhorn', 'sinkhorn_stabilized' or
+        'sinkhorn_epsilon_scaling', see those function for specific parameters
+    maxIter : int, optional
+        Max number of iterations
+    stopThr : float, optional
+        Stop threshol on error ( > 0 )
+    verbose : bool, optional
+        Print information along iterations
+    log : bool, optional
+        record log if True
+
+    Returns
+    -------
+    gamma : (na x nb) torch.tensor
+        Optimal transportation matrix for the given parameters
+    log : dict
+        log dictionary return only if log==True in parameters
+
+    References
+    ----------
+    [1] M. Cuturi, Sinkhorn Distances : Lightspeed Computation of Optimal Transport, Advances in Neural Information Processing Systems (NIPS) 26, 2013
+    See Also
+    --------
+
+    """
+
+    if method.lower() == 'sinkhorn':
+        return sinkhorn_knopp(a, b, C, reg, maxIter=maxIter,
+                              stopThr=stopThr, verbose=verbose, log=log,
+                              warm_start=warm_start, eval_freq=eval_freq, print_freq=print_freq,
+                              **kwargs)
+    elif method.lower() == 'sinkhorn_stabilized':
+        return sinkhorn_stabilized(a, b, C, reg, maxIter=maxIter, tau=tau,
+                                   stopThr=stopThr, verbose=verbose, log=log,
+                                   warm_start=warm_start, eval_freq=eval_freq, print_freq=print_freq,
+                                   **kwargs)
+    elif method.lower() == 'sinkhorn_epsilon_scaling':
+        return sinkhorn_epsilon_scaling(a, b, C, reg,
+                                        maxIter=maxIter, maxInnerIter=100, tau=tau,
+                                        scaling_base=0.75, scaling_coef=None, stopThr=stopThr,
+                                        verbose=False, log=log, warm_start=warm_start, eval_freq=eval_freq,
+                                        print_freq=print_freq, **kwargs)
+    else:
+        raise ValueError("Unknown method '%s'." % method)
+
+
+def sinkhorn_knopp(a, b, C, reg=1e-1, maxIter=1000, stopThr=1e-9,
+                   verbose=False, log=False, warm_start=None, eval_freq=10, print_freq=200, **kwargs):
+    """
+    Solve the entropic regularization optimal transport
+    The input should be PyTorch tensors
+    The function solves the following optimization problem:
+
+    .. math::
+        \gamma = arg\min_\gamma <\gamma,C>_F + reg\cdot\Omega(\gamma)
+        s.t. \gamma 1 = a
+             \gamma^T 1= b
+             \gamma\geq 0
+    where :
+    - C is the (ns,nt) metric cost matrix
+    - :math:`\Omega` is the entropic regularization term :math:`\Omega(\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})`
+    - a and b are target and source measures (sum to 1)
+    The algorithm used for solving the problem is the Sinkhorn-Knopp matrix scaling algorithm as proposed in [1].
+
+    Parameters
+    ----------
+    a : torch.tensor (na,)
+        samples measure in the target domain
+    b : torch.tensor (nb,)
+        samples in the source domain
+    C : torch.tensor (na,nb)
+        loss matrix
+    reg : float
+        Regularization term > 0
+    maxIter : int, optional
+        Max number of iterations
+    stopThr : float, optional
+        Stop threshol on error ( > 0 )
+    verbose : bool, optional
+        Print information along iterations
+    log : bool, optional
+        record log if True
+
+    Returns
+    -------
+    gamma : (na x nb) torch.tensor
+        Optimal transportation matrix for the given parameters
+    log : dict
+        log dictionary return only if log==True in parameters
+
+    References
+    ----------
+    [1] M. Cuturi, Sinkhorn Distances : Lightspeed Computation of Optimal Transport, Advances in Neural Information Processing Systems (NIPS) 26, 2013
+    See Also
+    --------
+
+    """
+
+    device = a.device
+    na, nb = C.shape
+
+    assert na >= 1 and nb >= 1, 'C needs to be 2d'
+    assert na == a.shape[0] and nb == b.shape[0], "Shape of a or b does't match that of C"
+    assert reg > 0, 'reg should be greater than 0'
+    assert a.min() >= 0. and b.min() >= 0., 'Elements in a or b less than 0'
+
+    if log:
+        log = {'err': []}
+
+    if warm_start is not None:
+        u = warm_start['u']
+        v = warm_start['v']
+    else:
+        u = torch.ones(na, dtype=a.dtype).to(device) / na
+        v = torch.ones(nb, dtype=b.dtype).to(device) / nb
+
+    K = torch.empty(C.shape, dtype=C.dtype).to(device)
+    torch.div(C, -reg, out=K)
+    torch.exp(K, out=K)
+
+    b_hat = torch.empty(b.shape, dtype=C.dtype).to(device)
+
+    it = 1
+    err = 1
+
+    # allocate memory beforehand
+    KTu = torch.empty(v.shape, dtype=v.dtype).to(device)
+    Kv = torch.empty(u.shape, dtype=u.dtype).to(device)
+
+    while (err > stopThr and it <= maxIter):
+        upre, vpre = u, v
+        torch.matmul(u, K, out=KTu)
+        v = torch.div(b, KTu + M_EPS)
+        torch.matmul(K, v, out=Kv)
+        u = torch.div(a, Kv + M_EPS)
+
+        if torch.any(torch.isnan(u)) or torch.any(torch.isnan(v)) or \
+                torch.any(torch.isinf(u)) or torch.any(torch.isinf(v)):
+            print('Warning: numerical errors at iteration', it)
+            u, v = upre, vpre
+            break
+
+        if log and it % eval_freq == 0:
+            # we can speed up the process by checking for the error only all
+            # the eval_freq iterations
+            # below is equivalent to:
+            # b_hat = torch.sum(u.reshape(-1, 1) * K * v.reshape(1, -1), 0)
+            # but with more memory efficient
+            b_hat = torch.matmul(u, K) * v
+            err = (b - b_hat).pow(2).sum().item()
+            # err = (b - b_hat).abs().sum().item()
+            log['err'].append(err)
+
+        if verbose and it % print_freq == 0:
+            print('iteration {:5d}, constraint error {:5e}'.format(it, err))
+
+        it += 1
+
+    if log:
+        log['u'] = u
+        log['v'] = v
+        log['alpha'] = reg * torch.log(u + M_EPS)
+        log['beta'] = reg * torch.log(v + M_EPS)
+
+    # transport plan
+    P = u.reshape(-1, 1) * K * v.reshape(1, -1)
+    if log:
+        return P, log
+    else:
+        return P
+
+
+def sinkhorn_stabilized(a, b, C, reg=1e-1, maxIter=1000, tau=1e3, stopThr=1e-9,
+                        verbose=False, log=False, warm_start=None, eval_freq=10, print_freq=200, **kwargs):
+    """
+    Solve the entropic regularization OT problem with log stabilization
+    The function solves the following optimization problem:
+
+    .. math::
+        \gamma = arg\min_\gamma <\gamma,C>_F + reg\cdot\Omega(\gamma)
+        s.t. \gamma 1 = a
+             \gamma^T 1= b
+             \gamma\geq 0
+    where :
+    - C is the (ns,nt) metric cost matrix
+    - :math:`\Omega` is the entropic regularization term :math:`\Omega(\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})`
+    - a and b are target and source measures (sum to 1)
+
+    The algorithm used for solving the problem is the Sinkhorn-Knopp matrix scaling algorithm as proposed in [1]
+    but with the log stabilization proposed in [3] an defined in [2] (Algo 3.1)
+
+    Parameters
+    ----------
+    a : torch.tensor (na,)
+        samples measure in the target domain
+    b : torch.tensor (nb,)
+        samples in the source domain
+    C : torch.tensor (na,nb)
+        loss matrix
+    reg : float
+        Regularization term > 0
+    tau : float
+        thershold for max value in u or v for log scaling
+    maxIter : int, optional
+        Max number of iterations
+    stopThr : float, optional
+        Stop threshol on error ( > 0 )
+    verbose : bool, optional
+        Print information along iterations
+    log : bool, optional
+        record log if True
+
+    Returns
+    -------
+    gamma : (na x nb) torch.tensor
+        Optimal transportation matrix for the given parameters
+    log : dict
+        log dictionary return only if log==True in parameters
+
+    References
+    ----------
+    [1] M. Cuturi, Sinkhorn Distances : Lightspeed Computation of Optimal Transport, Advances in Neural Information Processing Systems (NIPS) 26, 2013
+    [2] Bernhard Schmitzer. Stabilized Sparse Scaling Algorithms for Entropy Regularized Transport Problems. SIAM Journal on Scientific Computing, 2019
+    [3] Chizat, L., Peyré, G., Schmitzer, B., & Vialard, F. X. (2016). Scaling algorithms for unbalanced transport problems. arXiv preprint arXiv:1607.05816.
+
+    See Also
+    --------
+
+    """
+
+    device = a.device
+    na, nb = C.shape
+
+    assert na >= 1 and nb >= 1, 'C needs to be 2d'
+    assert na == a.shape[0] and nb == b.shape[0], "Shape of a or b does't match that of C"
+    assert reg > 0, 'reg should be greater than 0'
+    assert a.min() >= 0. and b.min() >= 0., 'Elements in a or b less than 0'
+
+    if log:
+        log = {'err': []}
+
+    if warm_start is not None:
+        alpha = warm_start['alpha']
+        beta = warm_start['beta']
+    else:
+        alpha = torch.zeros(na, dtype=a.dtype).to(device)
+        beta = torch.zeros(nb, dtype=b.dtype).to(device)
+
+    u = torch.ones(na, dtype=a.dtype).to(device) / na
+    v = torch.ones(nb, dtype=b.dtype).to(device) / nb
+
+    def update_K(alpha, beta):
+        """log space computation"""
+        """memory efficient"""
+        torch.add(alpha.reshape(-1, 1), beta.reshape(1, -1), out=K)
+        torch.add(K, -C, out=K)
+        torch.div(K, reg, out=K)
+        torch.exp(K, out=K)
+
+    def update_P(alpha, beta, u, v, ab_updated=False):
+        """log space P (gamma) computation"""
+        torch.add(alpha.reshape(-1, 1), beta.reshape(1, -1), out=P)
+        torch.add(P, -C, out=P)
+        torch.div(P, reg, out=P)
+        if not ab_updated:
+            torch.add(P, torch.log(u + M_EPS).reshape(-1, 1), out=P)
+            torch.add(P, torch.log(v + M_EPS).reshape(1, -1), out=P)
+        torch.exp(P, out=P)
+
+    K = torch.empty(C.shape, dtype=C.dtype).to(device)
+    update_K(alpha, beta)
+
+    b_hat = torch.empty(b.shape, dtype=C.dtype).to(device)
+
+    it = 1
+    err = 1
+    ab_updated = False
+
+    # allocate memory beforehand
+    KTu = torch.empty(v.shape, dtype=v.dtype).to(device)
+    Kv = torch.empty(u.shape, dtype=u.dtype).to(device)
+    P = torch.empty(C.shape, dtype=C.dtype).to(device)
+
+    while (err > stopThr and it <= maxIter):
+        upre, vpre = u, v
+        torch.matmul(u, K, out=KTu)
+        v = torch.div(b, KTu + M_EPS)
+        torch.matmul(K, v, out=Kv)
+        u = torch.div(a, Kv + M_EPS)
+
+        ab_updated = False
+        # remove numerical problems and store them in K
+        if u.abs().sum() > tau or v.abs().sum() > tau:
+            alpha += reg * torch.log(u + M_EPS)
+            beta += reg * torch.log(v + M_EPS)
+            u.fill_(1. / na)
+            v.fill_(1. / nb)
+            update_K(alpha, beta)
+            ab_updated = True
+
+        if log and it % eval_freq == 0:
+            # we can speed up the process by checking for the error only all
+            # the eval_freq iterations
+            update_P(alpha, beta, u, v, ab_updated)
+            b_hat = torch.sum(P, 0)
+            err = (b - b_hat).pow(2).sum().item()
+            log['err'].append(err)
+
+        if verbose and it % print_freq == 0:
+            print('iteration {:5d}, constraint error {:5e}'.format(it, err))
+
+        it += 1
+
+    if log:
+        log['u'] = u
+        log['v'] = v
+        log['alpha'] = alpha + reg * torch.log(u + M_EPS)
+        log['beta'] = beta + reg * torch.log(v + M_EPS)
+
+    # transport plan
+    update_P(alpha, beta, u, v, False)
+
+    if log:
+        return P, log
+    else:
+        return P
+
+
+def sinkhorn_epsilon_scaling(a, b, C, reg=1e-1, maxIter=100, maxInnerIter=100, tau=1e3, scaling_base=0.75,
+                             scaling_coef=None, stopThr=1e-9, verbose=False, log=False, warm_start=None, eval_freq=10,
+                             print_freq=200, **kwargs):
+    """
+    Solve the entropic regularization OT problem with log stabilization
+    The function solves the following optimization problem:
+
+    .. math::
+        \gamma = arg\min_\gamma <\gamma,C>_F + reg\cdot\Omega(\gamma)
+        s.t. \gamma 1 = a
+             \gamma^T 1= b
+             \gamma\geq 0
+    where :
+    - C is the (ns,nt) metric cost matrix
+    - :math:`\Omega` is the entropic regularization term :math:`\Omega(\gamma)=\sum_{i,j} \gamma_{i,j}\log(\gamma_{i,j})`
+    - a and b are target and source measures (sum to 1)
+
+    The algorithm used for solving the problem is the Sinkhorn-Knopp matrix
+    scaling algorithm as proposed in [1] but with the log stabilization
+    proposed in [3] and the log scaling proposed in [2] algorithm 3.2
+
+    Parameters
+    ----------
+    a : torch.tensor (na,)
+        samples measure in the target domain
+    b : torch.tensor (nb,)
+        samples in the source domain
+    C : torch.tensor (na,nb)
+        loss matrix
+    reg : float
+        Regularization term > 0
+    tau : float
+        thershold for max value in u or v for log scaling
+    maxIter : int, optional
+        Max number of iterations
+    stopThr : float, optional
+        Stop threshol on error ( > 0 )
+    verbose : bool, optional
+        Print information along iterations
+    log : bool, optional
+        record log if True
+
+    Returns
+    -------
+    gamma : (na x nb) torch.tensor
+        Optimal transportation matrix for the given parameters
+    log : dict
+        log dictionary return only if log==True in parameters
+
+    References
+    ----------
+    [1] M. Cuturi, Sinkhorn Distances : Lightspeed Computation of Optimal Transport, Advances in Neural Information Processing Systems (NIPS) 26, 2013
+    [2] Bernhard Schmitzer. Stabilized Sparse Scaling Algorithms for Entropy Regularized Transport Problems. SIAM Journal on Scientific Computing, 2019
+    [3] Chizat, L., Peyré, G., Schmitzer, B., & Vialard, F. X. (2016). Scaling algorithms for unbalanced transport problems. arXiv preprint arXiv:1607.05816.
+
+    See Also
+    --------
+
+    """
+
+    na, nb = C.shape
+
+    assert na >= 1 and nb >= 1, 'C needs to be 2d'
+    assert na == a.shape[0] and nb == b.shape[0], "Shape of a or b does't match that of C"
+    assert reg > 0, 'reg should be greater than 0'
+    assert a.min() >= 0. and b.min() >= 0., 'Elements in a or b less than 0'
+
+    def get_reg(it, reg, pre_reg):
+        if it == 1:
+            return scaling_coef
+        else:
+            if (pre_reg - reg) * scaling_base < M_EPS:
+                return reg
+            else:
+                return (pre_reg - reg) * scaling_base + reg
+
+    if scaling_coef is None:
+        scaling_coef = C.max() + reg
+
+    it = 1
+    err = 1
+    running_reg = scaling_coef
+
+    if log:
+        log = {'err': []}
+
+    warm_start = None
+
+    while (err > stopThr and it <= maxIter):
+        running_reg = get_reg(it, reg, running_reg)
+        P, _log = sinkhorn_stabilized(a, b, C, running_reg, maxIter=maxInnerIter, tau=tau,
+                                      stopThr=stopThr, verbose=False, log=True,
+                                      warm_start=warm_start, eval_freq=eval_freq, print_freq=print_freq,
+                                      **kwargs)
+
+        warm_start = {}
+        warm_start['alpha'] = _log['alpha']
+        warm_start['beta'] = _log['beta']
+
+        primal_val = (C * P).sum() + reg * (P * torch.log(P)).sum() - reg * P.sum()
+        dual_val = (_log['alpha'] * a).sum() + (_log['beta'] * b).sum() - reg * P.sum()
+        err = primal_val - dual_val
+        log['err'].append(err)
+
+        if verbose and it % print_freq == 0:
+            print('iteration {:5d}, constraint error {:5e}'.format(it, err))
+
+        it += 1
+
+    if log:
+        log['alpha'] = _log['alpha']
+        log['beta'] = _log['beta']
+        return P, log
+    else:
+        return P

code/losses/ot_loss.py

@@ -0,0 +1,67 @@
+import torch
+from torch.nn import Module
+from .bregman_pytorch import sinkhorn
+
+class OT_Loss(Module):
+    def __init__(self, c_size, stride, norm_cood, device, num_of_iter_in_ot=100, reg=10.0):
+        super(OT_Loss, self).__init__()
+        assert c_size % stride == 0
+
+        self.c_size = c_size
+        self.device = device
+        self.norm_cood = norm_cood
+        self.num_of_iter_in_ot = num_of_iter_in_ot
+        self.reg = reg
+
+        # coordinate is same to image space, set to constant since crop size is same
+        self.cood = torch.arange(0, c_size, step=stride,
+                                 dtype=torch.float32, device=device) + stride / 2
+        self.density_size = self.cood.size(0)
+        self.cood.unsqueeze_(0) # [1, #cood]
+        if self.norm_cood:
+            self.cood = self.cood / c_size * 2 - 1 # map to [-1, 1]
+        self.output_size = self.cood.size(1)
+
+
+    def forward(self, normed_density, unnormed_density, points):
+        batch_size = normed_density.size(0)
+        assert len(points) == batch_size
+        assert self.output_size == normed_density.size(2)
+        loss = torch.zeros([1]).to(self.device)
+        ot_obj_values = torch.zeros([1]).to(self.device)
+        wd = 0 # wasserstain distance
+        for idx, im_points in enumerate(points):
+            if len(im_points) > 0:
+                # compute l2 square distance, it should be source target distance. [#gt, #cood * #cood]
+                if self.norm_cood:
+                    im_points = im_points / self.c_size * 2 - 1 # map to [-1, 1]
+                x = im_points[:, 0].unsqueeze_(1)  # [#gt, 1]
+                y = im_points[:, 1].unsqueeze_(1)
+                x_dis = -2 * torch.matmul(x, self.cood) + x * x + self.cood * self.cood # [#gt, #cood]
+                y_dis = -2 * torch.matmul(y, self.cood) + y * y + self.cood * self.cood
+                y_dis.unsqueeze_(2)
+                x_dis.unsqueeze_(1)
+                dis = y_dis + x_dis
+                dis = dis.view((dis.size(0), -1)) # size of [#gt, #cood * #cood]
+
+                source_prob = normed_density[idx][0].view([-1]).detach()
+                target_prob = (torch.ones([len(im_points)]) / len(im_points)).to(self.device)
+                # use sinkhorn to solve OT, compute optimal beta.
+                P, log = sinkhorn(target_prob, source_prob, dis, self.reg, maxIter=self.num_of_iter_in_ot, log=True)
+                beta = log['beta'] # size is the same as source_prob: [#cood * #cood]
+                ot_obj_values += torch.sum(normed_density[idx] * beta.view([1, self.output_size, self.output_size]))
+                # compute the gradient of OT loss to predicted density (unnormed_density).
+                # im_grad = beta / source_count - < beta, source_density> / (source_count)^2
+                source_density = unnormed_density[idx][0].view([-1]).detach()
+                source_count = source_density.sum()
+                im_grad_1 = (source_count) / (source_count * source_count+1e-8) * beta # size of [#cood * #cood]
+                im_grad_2 = (source_density * beta).sum() / (source_count * source_count + 1e-8) # size of 1
+                im_grad = im_grad_1 - im_grad_2
+                im_grad = im_grad.detach().view([1, self.output_size, self.output_size])
+                # Define loss = <im_grad, predicted density>. The gradient of loss w.r.t prediced density is im_grad.
+                loss += torch.sum(unnormed_density[idx] * im_grad)
+                wd += torch.sum(dis * P).item()
+
+        return loss, wd, ot_obj_values
+
+

code/models.py

@@ -0,0 +1,57 @@
+import torch.nn as nn
+import torch.utils.model_zoo as model_zoo
+from torch.nn import functional as F
+
+__all__ = ['vgg19']
+model_urls = {
+    'vgg19': 'https://download.pytorch.org/models/vgg19-dcbb9e9d.pth',
+}
+
+class VGG(nn.Module):
+    def __init__(self, features):
+        super(VGG, self).__init__()
+        self.features = features
+        self.reg_layer = nn.Sequential(
+            nn.Conv2d(512, 256, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(256, 128, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True),
+        )
+        self.density_layer = nn.Sequential(nn.Conv2d(128, 1, 1), nn.ReLU())
+
+    def forward(self, x):
+        x = self.features(x)
+        x = F.upsample_bilinear(x, scale_factor=2)
+        x = self.reg_layer(x)
+        mu = self.density_layer(x)
+        B, C, H, W = mu.size()
+        mu_sum = mu.view([B, -1]).sum(1).unsqueeze(1).unsqueeze(2).unsqueeze(3)
+        mu_normed = mu / (mu_sum + 1e-6)
+        return mu, mu_normed
+
+def make_layers(cfg, batch_norm=False):
+    layers = []
+    in_channels = 3
+    for v in cfg:
+        if v == 'M':
+            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
+        else:
+            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
+            if batch_norm:
+                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
+            else:
+                layers += [conv2d, nn.ReLU(inplace=True)]
+            in_channels = v
+    return nn.Sequential(*layers)
+
+cfg = {
+    'E': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512]
+}
+
+def vgg19():
+    """VGG 19-layer model (configuration "E")
+        model pre-trained on ImageNet
+    """
+    model = VGG(make_layers(cfg['E']))
+    model.load_state_dict(model_zoo.load_url(model_urls['vgg19']), strict=False)
+    return model

code/preprocess/preprocess_dataset_nwpu.py

@@ -0,0 +1,137 @@
+from scipy.io import loadmat
+from PIL import Image
+import numpy as np
+import os
+import cv2
+
+
+def cal_new_size_v2(im_h, im_w, min_size, max_size):
+    rate = 1.0 * max_size / im_h
+    rate_w = im_w * rate
+    if rate_w > max_size:
+        rate = 1.0 * max_size / im_w
+    tmp_h = int(1.0 * im_h * rate / 16) * 16
+
+    if tmp_h < min_size:
+        rate = 1.0 * min_size / im_h
+    tmp_w = int(1.0 * im_w * rate / 16) * 16
+
+    if tmp_w < min_size:
+        rate = 1.0 * min_size / im_w
+    tmp_h = min(max(int(1.0 * im_h * rate / 16) * 16, min_size), max_size)
+    tmp_w = min(max(int(1.0 * im_w * rate / 16) * 16, min_size), max_size)
+
+    rate_h = 1.0 * tmp_h / im_h
+    rate_w = 1.0 * tmp_w / im_w
+    assert tmp_h >= min_size and tmp_h <= max_size
+    assert tmp_w >= min_size and tmp_w <= max_size
+    return tmp_h, tmp_w, rate_h, rate_w
+
+
+def gen_density_map_gaussian(im_height, im_width, points, sigma=4):
+    """
+    func: generate the density map.
+    points: [num_gt, 2], for each row: [width, height]
+    """
+    density_map = np.zeros([im_height, im_width], dtype=np.float32)
+    h, w = density_map.shape[:2]
+    num_gt = np.squeeze(points).shape[0]
+    if num_gt == 0:
+        return density_map
+    for p in points:
+        p = np.round(p).astype(int)
+        p[0], p[1] = min(h - 1, p[1]), min(w - 1, p[0])
+        gaussian_radius = sigma * 2 - 1
+        gaussian_map = np.multiply(
+            cv2.getGaussianKernel(int(gaussian_radius * 2 + 1), sigma),
+            cv2.getGaussianKernel(int(gaussian_radius * 2 + 1), sigma).T
+        )
+        x_left, x_right, y_up, y_down = 0, gaussian_map.shape[1], 0, gaussian_map.shape[0]
+        # cut the gaussian kernel
+        if p[1] < gaussian_radius:
+            x_left = gaussian_radius - p[1]
+        if p[0] < gaussian_radius:
+            y_up = gaussian_radius - p[0]
+        if p[1] + gaussian_radius >= w:
+            x_right = gaussian_map.shape[1] - (gaussian_radius + p[1] - w) - 1
+        if p[0] + gaussian_radius >= h:
+            y_down = gaussian_map.shape[0] - (gaussian_radius + p[0] - h) - 1
+        gaussian_map = gaussian_map[y_up:y_down, x_left:x_right]
+        if np.sum(gaussian_map):
+            gaussian_map = gaussian_map / np.sum(gaussian_map)
+        density_map[
+        max(0, p[0] - gaussian_radius):min(h, p[0] + gaussian_radius + 1),
+        max(0, p[1] - gaussian_radius):min(w, p[1] + gaussian_radius + 1)
+        ] += gaussian_map
+    density_map = density_map / (np.sum(density_map / num_gt))
+    return density_map
+
+
+def generate_data(im_path, mat_path, min_size, max_size):
+    im = Image.open(im_path).convert('RGB')
+    im_w, im_h = im.size
+    points = loadmat(mat_path)['annPoints'].astype(np.float32)
+    if len(points) > 0:  # some image has no crowd
+        idx_mask = (points[:, 0] >= 0) * (points[:, 0] <= im_w) * (points[:, 1] >= 0) * (points[:, 1] <= im_h)
+        points = points[idx_mask]
+    im_h, im_w, rr_h, rr_w = cal_new_size_v2(im_h, im_w, min_size, max_size)
+    im = np.array(im)
+    if rr_h != 1.0 or rr_w != 1.0:
+        im = cv2.resize(np.array(im), (im_w, im_h), cv2.INTER_CUBIC)
+        if len(points) > 0:  # some image has no crowd
+            points[:, 0] = points[:, 0] * rr_w
+            points[:, 1] = points[:, 1] * rr_h
+
+    density_map = gen_density_map_gaussian(im_h, im_w, points, sigma=8)
+    return Image.fromarray(im), points, density_map
+
+
+def generate_image(im_path, min_size, max_size):
+    im = Image.open(im_path).convert('RGB')
+    im_w, im_h = im.size
+    im_h, im_w, rr_h, rr_w = cal_new_size_v2(im_h, im_w, min_size, max_size)
+    im = np.array(im)
+    if rr_h != 1.0 or rr_w != 1.0:
+        im = cv2.resize(np.array(im), (im_w, im_h), cv2.INTER_CUBIC)
+    return Image.fromarray(im)
+
+
+def main(input_dataset_path, output_dataset_path, min_size=384, max_size=1920):
+    ori_img_path = os.path.join(input_dataset_path, 'images')
+    ori_anno_path = os.path.join(input_dataset_path, 'mats')
+
+    for phase in ['train', 'val']:
+        sub_save_dir = os.path.join(output_dataset_path, phase)
+        if not os.path.exists(sub_save_dir):
+            os.makedirs(sub_save_dir)
+        with open(os.path.join(input_dataset_path, '{}.txt'.format(phase))) as f:
+            lines = f.readlines()
+            for i in lines:
+                i = i.strip().split(' ')[0]
+                im_path = os.path.join(ori_img_path, i + '.jpg')
+                mat_path = os.path.join(ori_anno_path, i + '.mat')
+                name = os.path.basename(im_path)
+                im_save_path = os.path.join(sub_save_dir, name)
+                print(name)
+                # The Gaussian smoothed density map is just for visualization. It's not used in training.
+                im, points, density_map = generate_data(im_path, mat_path, min_size, max_size)
+                im.save(im_save_path)
+                gd_save_path = im_save_path.replace('jpg', 'npy')
+                np.save(gd_save_path, points)
+                dm_save_path = im_save_path.replace('.jpg', '_densitymap.npy')
+                np.save(dm_save_path, density_map)
+
+    for phase in ['test']:
+        sub_save_dir = os.path.join(output_dataset_path, phase)
+        if not os.path.exists(sub_save_dir):
+            os.makedirs(sub_save_dir)
+        with open(os.path.join(input_dataset_path, '{}.txt'.format(phase))) as f:
+            lines = f.readlines()
+            for i in lines:
+                i = i.strip().split(' ')[0]
+                im_path = os.path.join(ori_img_path, i + '.jpg')
+                name = os.path.basename(im_path)
+                im_save_path = os.path.join(sub_save_dir, name)
+                print(name)
+                im = generate_image(im_path, min_size, max_size)
+                im.save(im_save_path)

code/preprocess/preprocess_dataset_qnrf.py

@@ -0,0 +1,82 @@
+from scipy.io import loadmat
+from PIL import Image
+import numpy as np
+import os
+from glob import glob
+import cv2
+
+dir_name = os.path.dirname(os.path.abspath(__file__))
+
+def cal_new_size(im_h, im_w, min_size, max_size):
+    if im_h < im_w:
+        if im_h < min_size:
+            ratio = 1.0 * min_size / im_h
+            im_h = min_size
+            im_w = round(im_w * ratio)
+        elif im_h > max_size:
+            ratio = 1.0 * max_size / im_h
+            im_h = max_size
+            im_w = round(im_w * ratio)
+        else:
+            ratio = 1.0
+    else:
+        if im_w < min_size:
+            ratio = 1.0 * min_size / im_w
+            im_w = min_size
+            im_h = round(im_h * ratio)
+        elif im_w > max_size:
+            ratio = 1.0 * max_size / im_w
+            im_w = max_size
+            im_h = round(im_h * ratio)
+        else:
+            ratio = 1.0
+    return im_h, im_w, ratio
+
+
+def generate_data(im_path, min_size, max_size):
+    im = Image.open(im_path)
+    im_w, im_h = im.size
+    mat_path = im_path.replace('.jpg', '_ann.mat')
+    points = loadmat(mat_path)['annPoints'].astype(np.float32)
+    idx_mask = (points[:, 0] >= 0) * (points[:, 0] <= im_w) * (points[:, 1] >= 0) * (points[:, 1] <= im_h)
+    points = points[idx_mask]
+    im_h, im_w, rr = cal_new_size(im_h, im_w, min_size, max_size)
+    im = np.array(im)
+    if rr != 1.0:
+        im = cv2.resize(np.array(im), (im_w, im_h), cv2.INTER_CUBIC)
+        points = points * rr
+    return Image.fromarray(im), points
+
+
+def main(input_dataset_path, output_dataset_path, min_size=512, max_size=2048):
+    for phase in ['Train', 'Test']:
+        sub_dir = os.path.join(input_dataset_path, phase)
+        if phase == 'Train':
+            sub_phase_list = ['train', 'val']
+            for sub_phase in sub_phase_list:
+                sub_save_dir = os.path.join(output_dataset_path, sub_phase)
+                if not os.path.exists(sub_save_dir):
+                    os.makedirs(sub_save_dir)
+                with open(os.path.join(dir_name, 'qnrf_{}.txt'.format(sub_phase))) as f:
+                    for i in f:
+                        im_path = os.path.join(sub_dir, i.strip())
+                        name = os.path.basename(im_path)
+                        print(name)
+                        im, points = generate_data(im_path, min_size, max_size)
+                        im_save_path = os.path.join(sub_save_dir, name)
+                        im.save(im_save_path)
+                        gd_save_path = im_save_path.replace('jpg', 'npy')
+                        np.save(gd_save_path, points)
+        else:
+            sub_save_dir = os.path.join(output_dataset_path, 'test')
+            if not os.path.exists(sub_save_dir):
+                os.makedirs(sub_save_dir)
+            im_list = glob(os.path.join(sub_dir, '*jpg'))
+            for im_path in im_list:
+                name = os.path.basename(im_path)
+                print(name)
+                im, points = generate_data(im_path, min_size, max_size)
+                im_save_path = os.path.join(sub_save_dir, name)
+                im.save(im_save_path)
+                gd_save_path = im_save_path.replace('jpg', 'npy')
+                np.save(gd_save_path, points)

code/preprocess/qnrf_train.txt

@@ -0,0 +1,1081 @@
+img_0526.jpg
+img_0639.jpg
+img_0826.jpg
+img_0415.jpg
+img_0720.jpg
+img_0123.jpg
+img_0529.jpg
+img_1071.jpg
+img_0501.jpg
+img_0804.jpg
+img_0873.jpg
+img_0601.jpg
+img_0177.jpg
+img_0173.jpg
+img_0675.jpg
+img_1001.jpg
+img_0096.jpg
+img_1139.jpg
+img_0001.jpg
+img_0084.jpg
+img_0395.jpg
+img_0166.jpg
+img_0368.jpg
+img_0093.jpg
+img_0004.jpg
+img_0572.jpg
+img_0956.jpg
+img_0721.jpg
+img_0120.jpg
+img_0554.jpg
+img_0308.jpg
+img_0131.jpg
+img_0992.jpg
+img_0156.jpg
+img_0532.jpg
+img_0476.jpg
+img_0427.jpg
+img_1162.jpg
+img_0660.jpg
+img_0538.jpg
+img_0298.jpg
+img_0306.jpg
+img_1173.jpg
+img_1157.jpg
+img_0777.jpg
+img_0859.jpg
+img_0537.jpg
+img_0236.jpg
+img_0986.jpg
+img_0370.jpg
+img_0491.jpg
+img_1150.jpg
+img_0719.jpg
+img_1083.jpg
+img_0107.jpg
+img_1029.jpg
+img_0927.jpg
+img_0893.jpg
+img_0286.jpg
+img_1135.jpg
+img_0640.jpg
+img_0530.jpg
+img_1115.jpg
+img_0533.jpg
+img_0105.jpg
+img_0945.jpg
+img_1035.jpg
+img_0484.jpg
+img_1168.jpg
+img_0760.jpg
+img_0939.jpg
+img_0907.jpg
+img_0401.jpg
+img_0429.jpg
+img_0828.jpg
+img_1167.jpg
+img_0144.jpg
+img_0553.jpg
+img_0421.jpg
+img_0560.jpg
+img_0743.jpg
+img_0817.jpg
+img_0657.jpg
+img_0106.jpg
+img_0079.jpg
+img_0473.jpg
+img_0865.jpg
+img_0730.jpg
+img_0989.jpg
+img_0243.jpg
+img_0182.jpg
+img_0252.jpg
+img_0812.jpg
+img_0508.jpg
+img_0744.jpg
+img_0439.jpg
+img_0181.jpg
+img_0965.jpg
+img_0487.jpg
+img_0710.jpg
+img_1054.jpg
+img_0947.jpg
+img_0321.jpg
+img_0758.jpg
+img_0014.jpg
+img_0504.jpg
+img_0674.jpg
+img_0991.jpg
+img_0358.jpg
+img_1138.jpg
+img_0019.jpg
+img_0677.jpg
+img_0336.jpg
+img_0070.jpg
+img_0766.jpg
+img_0612.jpg
+img_1109.jpg
+img_0840.jpg
+img_0616.jpg
+img_0926.jpg
+img_0376.jpg
+img_0761.jpg
+img_0020.jpg
+img_0795.jpg
+img_0046.jpg
+img_0459.jpg
+img_0267.jpg
+img_0428.jpg
+img_1122.jpg
+img_0247.jpg
+img_1143.jpg
+img_0290.jpg
+img_0524.jpg
+img_0275.jpg
+img_1120.jpg
+img_0115.jpg
+img_0698.jpg
+img_0092.jpg
+img_0922.jpg
+img_1052.jpg
+img_0297.jpg
+img_0112.jpg
+img_0180.jpg
+img_0520.jpg
+img_0351.jpg
+img_0478.jpg
+img_0588.jpg
+img_0109.jpg
+img_0738.jpg
+img_0592.jpg
+img_0752.jpg
+img_1028.jpg
+img_1164.jpg
+img_0450.jpg
+img_0168.jpg
+img_1108.jpg
+img_0799.jpg
+img_0649.jpg
+img_0272.jpg
+img_0902.jpg
+img_0874.jpg
+img_0870.jpg
+img_0821.jpg
+img_0153.jpg
+img_0426.jpg
+img_0949.jpg
+img_0527.jpg
+img_1198.jpg
+img_0443.jpg
+img_0063.jpg
+img_0013.jpg
+img_0564.jpg
+img_0040.jpg
+img_0764.jpg
+img_0411.jpg
+img_0118.jpg
+img_1172.jpg
+img_0196.jpg
+img_0879.jpg
+img_0985.jpg
+img_0437.jpg
+img_0918.jpg
+img_0493.jpg
+img_0271.jpg
+img_0860.jpg
+img_0059.jpg
+img_0645.jpg
+img_1126.jpg
+img_0911.jpg
+img_1082.jpg
+img_0383.jpg
+img_0422.jpg
+img_0139.jpg
+img_1192.jpg
+img_0904.jpg
+img_0503.jpg
+img_0512.jpg
+img_0541.jpg
+img_0330.jpg
+img_0348.jpg
+img_0425.jpg
+img_0673.jpg
+img_0210.jpg
+img_0950.jpg
+img_0151.jpg
+img_0792.jpg
+img_0469.jpg
+img_0661.jpg
+img_0003.jpg
+img_0089.jpg
+img_0312.jpg
+img_0555.jpg
+img_0215.jpg
+img_0023.jpg
+img_1129.jpg
+img_0249.jpg
+img_0451.jpg
+img_1032.jpg
+img_0689.jpg
+img_1189.jpg
+img_0391.jpg
+img_0146.jpg
+img_0653.jpg
+img_0248.jpg
+img_0695.jpg
+img_0402.jpg
+img_0075.jpg
+img_1018.jpg
+img_1020.jpg
+img_0163.jpg
+img_0440.jpg
+img_0756.jpg
+img_0253.jpg
+img_0712.jpg
+img_0962.jpg
+img_0471.jpg
+img_0842.jpg
+img_0525.jpg
+img_1176.jpg
+img_1021.jpg
+img_0127.jpg
+img_0295.jpg
+img_1045.jpg
+img_1088.jpg
+img_1090.jpg
+img_0622.jpg
+img_0650.jpg
+img_0518.jpg
+img_0854.jpg
+img_0262.jpg
+img_0323.jpg
+img_0522.jpg
+img_0933.jpg
+img_0951.jpg
+img_0366.jpg
+img_0325.jpg
+img_1034.jpg
+img_0827.jpg
+img_0194.jpg
+img_0636.jpg
+img_0051.jpg
+img_0683.jpg
+img_0558.jpg
+img_0309.jpg
+img_0345.jpg
+img_0438.jpg
+img_1091.jpg
+img_0577.jpg
+img_0500.jpg
+img_0279.jpg
+img_1145.jpg
+img_0886.jpg
+img_1161.jpg
+img_0617.jpg
+img_0726.jpg
+img_0620.jpg
+img_0444.jpg
+img_1118.jpg
+img_0506.jpg
+img_0164.jpg
+img_0507.jpg
+img_0614.jpg
+img_0769.jpg
+img_1131.jpg
+img_0185.jpg
+img_0694.jpg
+img_1055.jpg
+img_0754.jpg
+img_0569.jpg
+img_0317.jpg
+img_0228.jpg
+img_0492.jpg
+img_1190.jpg
+img_0566.jpg
+img_0921.jpg
+img_0818.jpg
+img_0204.jpg
+img_0974.jpg
+img_0866.jpg
+img_1039.jpg
+img_0101.jpg
+img_0169.jpg
+img_0375.jpg
+img_0334.jpg
+img_1078.jpg
+img_0061.jpg
+img_0113.jpg
+img_0981.jpg
+img_0080.jpg
+img_0324.jpg
+img_0316.jpg
+img_0643.jpg
+img_0408.jpg
+img_0890.jpg
+img_0363.jpg
+img_0765.jpg
+img_0822.jpg
+img_0430.jpg
+img_0245.jpg
+img_0671.jpg
+img_0486.jpg
+img_1201.jpg
+img_0129.jpg
+img_1142.jpg
+img_0843.jpg
+img_1133.jpg
+img_0238.jpg
+img_0955.jpg
+img_1017.jpg
+img_0858.jpg
+img_1154.jpg
+img_0559.jpg
+img_0002.jpg
+img_0407.jpg
+img_1146.jpg
+img_1086.jpg
+img_0495.jpg
+img_0857.jpg
+img_0133.jpg
+img_0121.jpg
+img_0973.jpg
+img_0830.jpg
+img_0165.jpg
+img_0278.jpg
+img_1012.jpg
+img_0393.jpg
+img_0202.jpg
+img_0700.jpg
+img_0313.jpg
+img_0024.jpg
+img_0055.jpg
+img_0979.jpg
+img_0162.jpg
+img_0135.jpg
+img_0098.jpg
+img_0727.jpg
+img_0969.jpg
+img_1137.jpg
+img_0932.jpg
+img_1102.jpg
+img_0301.jpg
+img_0047.jpg
+img_0595.jpg
+img_0805.jpg
+img_0801.jpg
+img_1151.jpg
+img_0387.jpg
+img_0999.jpg
+img_0136.jpg
+img_1037.jpg
+img_1087.jpg
+img_1186.jpg
+img_0032.jpg
+img_0195.jpg
+img_0360.jpg
+img_0276.jpg
+img_0642.jpg
+img_0913.jpg
+img_0231.jpg
+img_0670.jpg
+img_1123.jpg
+img_0517.jpg
+img_0707.jpg
+img_0088.jpg
+img_0594.jpg
+img_0838.jpg
+img_0848.jpg
+img_0354.jpg
+img_0936.jpg
+img_0876.jpg
+img_1081.jpg
+img_0322.jpg
+img_0637.jpg
+img_0739.jpg
+img_0917.jpg
+img_0244.jpg
+img_0591.jpg
+img_0628.jpg
+img_0964.jpg
+img_0691.jpg
+img_0609.jpg
+img_0342.jpg
+img_1097.jpg
+img_1077.jpg
+img_0502.jpg
+img_0423.jpg
+img_0561.jpg
+img_1059.jpg
+img_0568.jpg
+img_0920.jpg
+img_0389.jpg
+img_0940.jpg
+img_0787.jpg
+img_0634.jpg
+img_0516.jpg
+img_0900.jpg
+img_0463.jpg
+img_0942.jpg
+img_0796.jpg
+img_0835.jpg
+img_0789.jpg
+img_0184.jpg
+img_0397.jpg
+img_1195.jpg
+img_1089.jpg
+img_0319.jpg
+img_0328.jpg
+img_0724.jpg
+img_0852.jpg
+img_0662.jpg
+img_0225.jpg
+img_0479.jpg
+img_0266.jpg
+img_0499.jpg
+img_0134.jpg
+img_1023.jpg
+img_1064.jpg
+img_0400.jpg
+img_0226.jpg
+img_0015.jpg
+img_0203.jpg
+img_0548.jpg
+img_1084.jpg
+img_0970.jpg
+img_0718.jpg
+img_0138.jpg
+img_0095.jpg
+img_0831.jpg
+img_0482.jpg
+img_1000.jpg
+img_0234.jpg
+img_0183.jpg
+img_0687.jpg
+img_0923.jpg
+img_0197.jpg
+img_1016.jpg
+img_1100.jpg
+img_0034.jpg
+img_0587.jpg
+img_0229.jpg
+img_1178.jpg
+img_0124.jpg
+img_0424.jpg
+img_0496.jpg
+img_0179.jpg
+img_1110.jpg
+img_0998.jpg
+img_0742.jpg
+img_0578.jpg
+img_0207.jpg
+img_0305.jpg
+img_0373.jpg
+img_0971.jpg
+img_0292.jpg
+img_0861.jpg
+img_0621.jpg
+img_0414.jpg
+img_1140.jpg
+img_0737.jpg
+img_0176.jpg
+img_1057.jpg
+img_1095.jpg
+img_0667.jpg
+img_0755.jpg
+img_0318.jpg
+img_0170.jpg
+img_0418.jpg
+img_0178.jpg
+img_1200.jpg
+img_0021.jpg
+img_0652.jpg
+img_0327.jpg
+img_0627.jpg
+img_1051.jpg
+img_0837.jpg
+img_0352.jpg
+img_0029.jpg
+img_0833.jpg
+img_0952.jpg
+img_0488.jpg
+img_0474.jpg
+img_0702.jpg
+img_0819.jpg
+img_1188.jpg
+img_0261.jpg
+img_0685.jpg
+img_1024.jpg
+img_0008.jpg
+img_0734.jpg
+img_0509.jpg
+img_0888.jpg
+img_0676.jpg
+img_0404.jpg
+img_1046.jpg
+img_1127.jpg
+img_1008.jpg
+img_0161.jpg
+img_0699.jpg
+img_0085.jpg
+img_0703.jpg
+img_0083.jpg
+img_0934.jpg
+img_0626.jpg
+img_1170.jpg
+img_1065.jpg
+img_0664.jpg
+img_0883.jpg
+img_0655.jpg
+img_0263.jpg
+img_1005.jpg
+img_1061.jpg
+img_0333.jpg
+img_0881.jpg
+img_1041.jpg
+img_0540.jpg
+img_1185.jpg
+img_0953.jpg
+img_0586.jpg
+img_1011.jpg
+img_0846.jpg
+img_0149.jpg
+img_1075.jpg
+img_0894.jpg
+img_0759.jpg
+img_1177.jpg
+img_0258.jpg
+img_0171.jpg
+img_0740.jpg
+img_0006.jpg
+img_0353.jpg
+img_0615.jpg
+img_0810.jpg
+img_0142.jpg
+img_0958.jpg
+img_0584.jpg
+img_0390.jpg
+img_0585.jpg
+img_0365.jpg
+img_0026.jpg
+img_0458.jpg
+img_0143.jpg
+img_0575.jpg
+img_1027.jpg
+img_1183.jpg
+img_0535.jpg
+img_0891.jpg
+img_1085.jpg
+img_0757.jpg
+img_0549.jpg
+img_0436.jpg
+img_0815.jpg
+img_0635.jpg
+img_0954.jpg
+img_0367.jpg
+img_0064.jpg
+img_0410.jpg
+img_0277.jpg
+img_1111.jpg
+img_1025.jpg
+img_0434.jpg
+img_1175.jpg
+img_1171.jpg
+img_0610.jpg
+img_0618.jpg
+img_0208.jpg
+img_0281.jpg
+img_0058.jpg
+img_0851.jpg
+img_0300.jpg
+img_0017.jpg
+img_0110.jpg
+img_0265.jpg
+img_0362.jpg
+img_1038.jpg
+img_0580.jpg
+img_1096.jpg
+img_0972.jpg
+img_0666.jpg
+img_0090.jpg
+img_1007.jpg
+img_0982.jpg
+img_0287.jpg
+img_0714.jpg
+img_0218.jpg
+img_0832.jpg
+img_0145.jpg
+img_0072.jpg
+img_0222.jpg
+img_0137.jpg
+img_0741.jpg
+img_0028.jpg
+img_0413.jpg
+img_0232.jpg
+img_0573.jpg
+img_0849.jpg
+img_0855.jpg
+img_0770.jpg
+img_0283.jpg
+img_0914.jpg
+img_0611.jpg
+img_1047.jpg
+img_0596.jpg
+img_0706.jpg
+img_0847.jpg
+img_0868.jpg
+img_0193.jpg
+img_0780.jpg
+img_0100.jpg
+img_0786.jpg
+img_0337.jpg
+img_0728.jpg
+img_0656.jpg
+img_0602.jpg
+img_1015.jpg
+img_0273.jpg
+img_0797.jpg
+img_0398.jpg
+img_0693.jpg
+img_0944.jpg
+img_0593.jpg
+img_0768.jpg
+img_0995.jpg
+img_1125.jpg
+img_0078.jpg
+img_0543.jpg
+img_0167.jpg
+img_0420.jpg
+img_0264.jpg
+img_0016.jpg
+img_0599.jpg
+img_0417.jpg
+img_0448.jpg
+img_0748.jpg
+img_0311.jpg
+img_0071.jpg
+img_0749.jpg
+img_0941.jpg
+img_0237.jpg
+img_0214.jpg
+img_1149.jpg
+img_0241.jpg
+img_0461.jpg
+img_0018.jpg
+img_0356.jpg
+img_0483.jpg
+img_0099.jpg
+img_0130.jpg
+img_0372.jpg
+img_0800.jpg
+img_0654.jpg
+img_0544.jpg
+img_1099.jpg
+img_1068.jpg
+img_0326.jpg
+img_0374.jpg
+img_0074.jpg
+img_0938.jpg
+img_0117.jpg
+img_0456.jpg
+img_0901.jpg
+img_0713.jpg
+img_0788.jpg
+img_0665.jpg
+img_0294.jpg
+img_0841.jpg
+img_0269.jpg
+img_0579.jpg
+img_1098.jpg
+img_0466.jpg
+img_0480.jpg
+img_0709.jpg
+img_0672.jpg
+img_1010.jpg
+img_0314.jpg
+img_0043.jpg
+img_0349.jpg
+img_0172.jpg
+img_1187.jpg
+img_0371.jpg
+img_0320.jpg
+img_1103.jpg
+img_1159.jpg
+img_0629.jpg
+img_0399.jpg
+img_0663.jpg
+img_0335.jpg
+img_1148.jpg
+img_0108.jpg
+img_0254.jpg
+img_0432.jpg
+img_0915.jpg
+img_0624.jpg
+img_0997.jpg
+img_0711.jpg
+img_0704.jpg
+img_1147.jpg
+img_0036.jpg
+img_0519.jpg
+img_0680.jpg
+img_0498.jpg
+img_0651.jpg
+img_0230.jpg
+img_0198.jpg
+img_0905.jpg
+img_0751.jpg
+img_0928.jpg
+img_0630.jpg
+img_0140.jpg
+img_0644.jpg
+img_0776.jpg
+img_0057.jpg
+img_0361.jpg
+img_0209.jpg
+img_0158.jpg
+img_1160.jpg
+img_1169.jpg
+img_0735.jpg
+img_0551.jpg
+img_0681.jpg
+img_0515.jpg
+img_0077.jpg
+img_0968.jpg
+img_0240.jpg
+img_1166.jpg
+img_0937.jpg
+img_0877.jpg
+img_0513.jpg
+img_0528.jpg
+img_0150.jpg
+img_1165.jpg
+img_0200.jpg
+img_0246.jpg
+img_0869.jpg
+img_0011.jpg
+img_0160.jpg
+img_0464.jpg
+img_0285.jpg
+img_0132.jpg
+img_0701.jpg
+img_0082.jpg
+img_1182.jpg
+img_0030.jpg
+img_0126.jpg
+img_0632.jpg
+img_0731.jpg
+img_0875.jpg
+img_0978.jpg
+img_0717.jpg
+img_0460.jpg
+img_1044.jpg
+img_1194.jpg
+img_0910.jpg
+img_0049.jpg
+img_0331.jpg
+img_0213.jpg
+img_0885.jpg
+img_0468.jpg
+img_0419.jpg
+img_1158.jpg
+img_0022.jpg
+img_0174.jpg
+img_0747.jpg
+img_1006.jpg
+img_0381.jpg
+img_1036.jpg
+img_0863.jpg
+img_0994.jpg
+img_0783.jpg
+img_0346.jpg
+img_0233.jpg
+img_0820.jpg
+img_1107.jpg
+img_1193.jpg
+img_0943.jpg
+img_1191.jpg
+img_0005.jpg
+img_0087.jpg
+img_0039.jpg
+img_0813.jpg
+img_0239.jpg
+img_0206.jpg
+img_0256.jpg
+img_1070.jpg
+img_0409.jpg
+img_0377.jpg
+img_0446.jpg
+img_0216.jpg
+img_0189.jpg
+img_0785.jpg
+img_0041.jpg
+img_0598.jpg
+img_0310.jpg
+img_0307.jpg
+img_1093.jpg
+img_0465.jpg
+img_0746.jpg
+img_0380.jpg
+img_0732.jpg
+img_0781.jpg
+img_0906.jpg
+img_0619.jpg
+img_0604.jpg
+img_0983.jpg
+img_0753.jpg
+img_0211.jpg
+img_0552.jpg
+img_0892.jpg
+img_0767.jpg
+img_1180.jpg
+img_1069.jpg
+img_0154.jpg
+img_0899.jpg
+img_0343.jpg
+img_0025.jpg
+img_1196.jpg
+img_0155.jpg
+img_0433.jpg
+img_0597.jpg
+img_0570.jpg
+img_0867.jpg
+img_0223.jpg
+img_0581.jpg
+img_0186.jpg
+img_0122.jpg
+img_1134.jpg
+img_0340.jpg
+img_0957.jpg
+img_0364.jpg
+img_0069.jpg
+img_1114.jpg
+img_0646.jpg
+img_0679.jpg
+img_0623.jpg
+img_0392.jpg
+img_0814.jpg
+img_0589.jpg
+img_0299.jpg
+img_0931.jpg
+img_0836.jpg
+img_0963.jpg
+img_0094.jpg
+img_0987.jpg
+img_0930.jpg
+img_0976.jpg
+img_0924.jpg
+img_0384.jpg
+img_0035.jpg
+img_0076.jpg
+img_1101.jpg
+img_0405.jpg
+img_0350.jpg
+img_0147.jpg
+img_0659.jpg
+img_1013.jpg
+img_0948.jpg
+img_0066.jpg
+img_1132.jpg
+img_0829.jpg
+img_0690.jpg
+img_1060.jpg
+img_0457.jpg
+img_0897.jpg
+img_0825.jpg
+img_1163.jpg
+img_0803.jpg
+img_0563.jpg
+img_0574.jpg
+img_0175.jpg
+img_1112.jpg
+img_0668.jpg
+img_0045.jpg
+img_0259.jpg
+img_0341.jpg
+img_1067.jpg
+img_1040.jpg
+img_1106.jpg
+img_0205.jpg
+img_0296.jpg
+img_0255.jpg
+img_1152.jpg
+img_0772.jpg
+img_0613.jpg
+img_1121.jpg
+img_0834.jpg
+img_0406.jpg
+img_0762.jpg
+img_0442.jpg
+img_0192.jpg
+img_0044.jpg
+img_0774.jpg
+img_0606.jpg
+img_0359.jpg
+img_0467.jpg
+img_0779.jpg
+img_0060.jpg
+img_1074.jpg
+img_0494.jpg
+img_1153.jpg
+img_0102.jpg
+img_0582.jpg
+img_0386.jpg
+img_0212.jpg
+img_0625.jpg
+img_0844.jpg
+img_0872.jpg
+img_1105.jpg
+img_0396.jpg
+img_1119.jpg
+img_0052.jpg
+img_0454.jpg
+img_1179.jpg
+img_0862.jpg
+img_0481.jpg
+img_1026.jpg
+img_0511.jpg
+img_0912.jpg
+img_1124.jpg
+img_0148.jpg
+img_0960.jpg
+img_0523.jpg
+img_0531.jpg
+img_0729.jpg
+img_0571.jpg
+img_0908.jpg
+img_0889.jpg
+img_0188.jpg
+img_0037.jpg
+img_0716.jpg
+img_1014.jpg
+img_0394.jpg
+img_1056.jpg
+img_0462.jpg
+img_0850.jpg
+img_0784.jpg
+img_1002.jpg
+img_0763.jpg
+img_0159.jpg
+img_0009.jpg
+img_0708.jpg
+img_1050.jpg
+img_0678.jpg
+img_0648.jpg
+img_0010.jpg
+img_1031.jpg
+img_0445.jpg
+img_0355.jpg
+img_1117.jpg
+img_0378.jpg
+img_0550.jpg
+img_0217.jpg
+img_0260.jpg
+img_0816.jpg
+img_0996.jpg
+img_0081.jpg
+img_0878.jpg
+img_0199.jpg
+img_0431.jpg
+img_1144.jpg
+img_0688.jpg
+img_0745.jpg
+img_0686.jpg
+img_1042.jpg
+img_0187.jpg
+img_1066.jpg
+img_0682.jpg
+img_0048.jpg
+img_0896.jpg
+img_0608.jpg
+img_1003.jpg
+img_1156.jpg
+img_0723.jpg
+img_0692.jpg
+img_0220.jpg
+img_0993.jpg
+img_1197.jpg
+img_0447.jpg
+img_0369.jpg
+img_0056.jpg
+img_0807.jpg
+img_0315.jpg
+img_0567.jpg
+img_0452.jpg
+img_1128.jpg
+img_0647.jpg
+img_0242.jpg
+img_0201.jpg
+img_0497.jpg
+img_0031.jpg
+img_0771.jpg
+img_0547.jpg
+img_0705.jpg
+img_0725.jpg
+img_1058.jpg
+img_0053.jpg
+img_1043.jpg
+img_0722.jpg
+img_0435.jpg
+img_0284.jpg
+img_0583.jpg
+img_0882.jpg
+img_0111.jpg
+img_0959.jpg
+img_1076.jpg
+img_0880.jpg
+img_0224.jpg
+img_0977.jpg
+img_0270.jpg
+img_0793.jpg
+img_0603.jpg
+img_1116.jpg
+img_0304.jpg
+img_0884.jpg
+img_1136.jpg
+img_0235.jpg
+img_0412.jpg
+img_0980.jpg
+img_0988.jpg
+img_0773.jpg
+img_1174.jpg
+img_0562.jpg
+img_0871.jpg
+img_0798.jpg
+img_0453.jpg
+img_0696.jpg
+img_0104.jpg
+img_0607.jpg
+img_0669.jpg
+img_0293.jpg
+img_1141.jpg
+img_0329.jpg
+img_0534.jpg
+img_1113.jpg
+img_0288.jpg
+img_0961.jpg
+img_0388.jpg
+img_0073.jpg
+img_0141.jpg
+img_0935.jpg
+img_1062.jpg
+img_0227.jpg
+img_0895.jpg
+img_0449.jpg
+img_0565.jpg
+img_1009.jpg
+img_0282.jpg
+img_0806.jpg
+img_1033.jpg
+img_0332.jpg
+img_0903.jpg
+img_0475.jpg
+img_0050.jpg
+img_0455.jpg
+img_0845.jpg
+img_0946.jpg
+img_0490.jpg
+img_0274.jpg
+img_0909.jpg
+img_0966.jpg
+img_0219.jpg
+img_0898.jpg
+img_0403.jpg

code/preprocess/qnrf_val.txt

@@ -0,0 +1,120 @@
+img_0042.jpg
+img_0697.jpg
+img_0012.jpg
+img_0062.jpg
+img_0990.jpg
+img_1048.jpg
+img_0576.jpg
+img_0802.jpg
+img_0116.jpg
+img_0119.jpg
+img_0967.jpg
+img_0054.jpg
+img_0782.jpg
+img_0514.jpg
+img_0929.jpg
+img_0809.jpg
+img_0033.jpg
+img_0125.jpg
+img_0633.jpg
+img_0038.jpg
+img_0775.jpg
+img_0600.jpg
+img_0157.jpg
+img_0824.jpg
+img_0103.jpg
+img_0984.jpg
+img_0250.jpg
+img_0505.jpg
+img_0631.jpg
+img_0556.jpg
+img_1049.jpg
+img_1181.jpg
+img_0097.jpg
+img_0536.jpg
+img_1104.jpg
+img_0733.jpg
+img_1130.jpg
+img_0808.jpg
+img_0086.jpg
+img_0302.jpg
+img_0114.jpg
+img_0470.jpg
+img_0715.jpg
+img_0641.jpg
+img_0557.jpg
+img_0510.jpg
+img_0152.jpg
+img_0485.jpg
+img_0190.jpg
+img_0065.jpg
+img_0839.jpg
+img_0068.jpg
+img_0864.jpg
+img_0477.jpg
+img_0441.jpg
+img_0546.jpg
+img_0091.jpg
+img_0853.jpg
+img_0975.jpg
+img_0357.jpg
+img_1004.jpg
+img_0794.jpg
+img_0750.jpg
+img_0791.jpg
+img_0605.jpg
+img_0590.jpg
+img_0489.jpg
+img_0191.jpg
+img_0007.jpg
+img_0778.jpg
+img_0658.jpg
+img_0289.jpg
+img_0925.jpg
+img_1184.jpg
+img_0521.jpg
+img_0291.jpg
+img_0823.jpg
+img_0382.jpg
+img_0416.jpg
+img_0736.jpg
+img_0268.jpg
+img_0128.jpg
+img_0280.jpg
+img_1022.jpg
+img_0545.jpg
+img_0257.jpg
+img_0251.jpg
+img_0684.jpg
+img_1092.jpg
+img_0638.jpg
+img_1079.jpg
+img_0790.jpg
+img_0811.jpg
+img_0303.jpg
+img_0542.jpg
+img_1019.jpg
+img_0472.jpg
+img_0027.jpg
+img_0539.jpg
+img_0856.jpg
+img_1094.jpg
+img_1030.jpg
+img_1063.jpg
+img_0887.jpg
+img_0067.jpg
+img_0379.jpg
+img_0919.jpg
+img_1155.jpg
+img_0221.jpg
+img_1053.jpg
+img_0916.jpg
+img_1072.jpg
+img_0347.jpg
+img_1199.jpg
+img_1080.jpg
+img_0385.jpg
+img_0344.jpg
+img_1073.jpg
+img_0339.jpg
+img_0338.jpg

code/preprocess_dataset.py

@@ -0,0 +1,23 @@
+# Preprocess images in QNRF and NWPU dataset.
+
+import argparse
+
+parser = argparse.ArgumentParser(description='Preprocess')
+parser.add_argument('--dataset', default='qnrf',
+                    help='dataset name, only support qnrf and nwpu')
+parser.add_argument('--input-dataset-path', default='data/QNRF',
+                    help='original data directory')
+parser.add_argument('--output-dataset-path', default='data/QNRF-Train-Val-Test',
+                    help='processed data directory')
+args = parser.parse_args()
+
+if args.dataset.lower() == 'qnrf':
+    from preprocess.preprocess_dataset_qnrf import main
+
+    main(args.input_dataset_path, args.output_dataset_path, 512, 2048)
+elif args.dataset.lower() == 'nwpu':
+    from preprocess.preprocess_dataset_nwpu import main
+
+    main(args.input_dataset_path, args.output_dataset_path, 384, 1920)
+else:
+    raise NotImplementedError

code/pretrained_models/model_nwpu.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f2a0c92ac22b5c6aee08b59ad9f002561e75da07555b58c20dd699db8aac59b2
+size 86005202

code/pretrained_models/model_qnrf.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:16ef954a2cef40c66ee664f69273559553b735d5e2c9f90e2444f3c25dd45e05
+size 86005202

code/requirements.txt

@@ -0,0 +1,8 @@
+torch
+torchvision
+numpy>=1.16.5
+scipy>=1.3.0
+opencv-python
+gdown
+Pillow
+gradio

code/test.py

@@ -0,0 +1,73 @@
+import argparse
+import torch
+import os
+import numpy as np
+import datasets.crowd as crowd
+from models import vgg19
+
+parser = argparse.ArgumentParser(description='Test ')
+parser.add_argument('--device', default='0', help='assign device')
+parser.add_argument('--crop-size', type=int, default=512,
+                    help='the crop size of the train image')
+parser.add_argument('--model-path', type=str, default='pretrained_models/model_qnrf.pth',
+                    help='saved model path')
+parser.add_argument('--data-path', type=str,
+                    default='data/QNRF-Train-Val-Test',
+                    help='saved model path')
+parser.add_argument('--dataset', type=str, default='qnrf',
+                    help='dataset name: qnrf, nwpu, sha, shb')
+parser.add_argument('--pred-density-map-path', type=str, default='',
+                    help='save predicted density maps when pred-density-map-path is not empty.')
+
+
+args = parser.parse_args()
+
+os.environ['CUDA_VISIBLE_DEVICES'] = args.device  # set vis gpu
+device = torch.device('cuda')
+
+model_path = args.model_path
+crop_size = args.crop_size
+data_path = args.data_path
+if args.dataset.lower() == 'qnrf':
+    dataset = crowd.Crowd_qnrf(os.path.join(data_path, 'test'), crop_size, 8, method='val')
+elif args.dataset.lower() == 'nwpu':
+    dataset = crowd.Crowd_nwpu(os.path.join(data_path, 'val'), crop_size, 8, method='val')
+elif args.dataset.lower() == 'sha' or args.dataset.lower() == 'shb':
+    dataset = crowd.Crowd_sh(os.path.join(data_path, 'test_data'), crop_size, 8, method='val')
+else:
+    raise NotImplementedError
+dataloader = torch.utils.data.DataLoader(dataset, 1, shuffle=False,
+                                         num_workers=1, pin_memory=True)
+
+if args.pred_density_map_path:
+    import cv2
+    if not os.path.exists(args.pred_density_map_path):
+        os.makedirs(args.pred_density_map_path)
+
+model = vgg19()
+model.to(device)
+model.load_state_dict(torch.load(model_path, device))
+model.eval()
+image_errs = []
+for inputs, count, name in dataloader:
+    inputs = inputs.to(device)
+    assert inputs.size(0) == 1, 'the batch size should equal to 1'
+    with torch.set_grad_enabled(False):
+        outputs, _ = model(inputs)
+    img_err = count[0].item() - torch.sum(outputs).item()
+
+    print(name, img_err, count[0].item(), torch.sum(outputs).item())
+    image_errs.append(img_err)
+
+    if args.pred_density_map_path:
+        vis_img = outputs[0, 0].cpu().numpy()
+        # normalize density map values from 0 to 1, then map it to 0-255.
+        vis_img = (vis_img - vis_img.min()) / (vis_img.max() - vis_img.min() + 1e-5)
+        vis_img = (vis_img * 255).astype(np.uint8)
+        vis_img = cv2.applyColorMap(vis_img, cv2.COLORMAP_JET)
+        cv2.imwrite(os.path.join(args.pred_density_map_path, str(name[0]) + '.png'), vis_img)
+
+image_errs = np.array(image_errs)
+mse = np.sqrt(np.mean(np.square(image_errs)))
+mae = np.mean(np.abs(image_errs))
+print('{}: mae {}, mse {}\n'.format(model_path, mae, mse))

code/train.py

@@ -0,0 +1,64 @@
+import argparse
+import os
+import torch
+from train_helper import Trainer
+
+
+def str2bool(v):
+    return v.lower() in ("yes", "true", "t", "1")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Train')
+    parser.add_argument('--data-dir', default='data/UCF-Train-Val-Test', help='data path')
+    parser.add_argument('--dataset', default='qnrf', help='dataset name: qnrf, nwpu, sha, shb')
+    parser.add_argument('--lr', type=float, default=1e-5,
+                        help='the initial learning rate')
+    parser.add_argument('--weight-decay', type=float, default=1e-4,
+                        help='the weight decay')
+    parser.add_argument('--resume', default='', type=str,
+                        help='the path of resume training model')
+    parser.add_argument('--max-epoch', type=int, default=1000,
+                        help='max training epoch')
+    parser.add_argument('--val-epoch', type=int, default=5,
+                        help='the num of steps to log training information')
+    parser.add_argument('--val-start', type=int, default=50,
+                        help='the epoch start to val')
+    parser.add_argument('--batch-size', type=int, default=10,
+                        help='train batch size')
+    parser.add_argument('--device', default='0', help='assign device')
+    parser.add_argument('--num-workers', type=int, default=3,
+                        help='the num of training process')
+    parser.add_argument('--crop-size', type=int, default=512,
+                        help='the crop size of the train image')
+    parser.add_argument('--wot', type=float, default=0.1, help='weight on OT loss')
+    parser.add_argument('--wtv', type=float, default=0.01, help='weight on TV loss')
+    parser.add_argument('--reg', type=float, default=10.0,
+                        help='entropy regularization in sinkhorn')
+    parser.add_argument('--num-of-iter-in-ot', type=int, default=100,
+                        help='sinkhorn iterations')
+    parser.add_argument('--norm-cood', type=int, default=0, help='whether to norm cood when computing distance')
+
+    args = parser.parse_args()
+
+    if args.dataset.lower() == 'qnrf':
+        args.crop_size = 512
+    elif args.dataset.lower() == 'nwpu':
+        args.crop_size = 384
+        args.val_epoch = 50
+    elif args.dataset.lower() == 'sha':
+        args.crop_size = 256
+    elif args.dataset.lower() == 'shb':
+        args.crop_size = 512
+    else:
+        raise NotImplementedError
+    return args
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    torch.backends.cudnn.benchmark = True
+    os.environ['CUDA_VISIBLE_DEVICES'] = args.device.strip()  # set vis gpu
+    trainer = Trainer(args)
+    trainer.setup()
+    trainer.train()

code/train_helper.py

@@ -0,0 +1,211 @@
+import os
+import time
+import torch
+import torch.nn as nn
+from torch import optim
+from torch.utils.data import DataLoader
+from torch.utils.data.dataloader import default_collate
+import numpy as np
+from datetime import datetime
+
+from datasets.crowd import Crowd_qnrf, Crowd_nwpu, Crowd_sh
+from models import vgg19
+from losses.ot_loss import OT_Loss
+from utils.pytorch_utils import Save_Handle, AverageMeter
+import utils.log_utils as log_utils
+
+
+def train_collate(batch):
+    transposed_batch = list(zip(*batch))
+    images = torch.stack(transposed_batch[0], 0)
+    points = transposed_batch[1]  # the number of points is not fixed, keep it as a list of tensor
+    gt_discretes = torch.stack(transposed_batch[2], 0)
+    return images, points, gt_discretes
+
+
+class Trainer(object):
+    def __init__(self, args):
+        self.args = args
+
+    def setup(self):
+        args = self.args
+        sub_dir = 'input-{}_wot-{}_wtv-{}_reg-{}_nIter-{}_normCood-{}'.format(
+            args.crop_size, args.wot, args.wtv, args.reg, args.num_of_iter_in_ot, args.norm_cood)
+
+        self.save_dir = os.path.join('ckpts', sub_dir)
+        if not os.path.exists(self.save_dir):
+            os.makedirs(self.save_dir)
+
+        time_str = datetime.strftime(datetime.now(), '%m%d-%H%M%S')
+        self.logger = log_utils.get_logger(os.path.join(self.save_dir, 'train-{:s}.log'.format(time_str)))
+        log_utils.print_config(vars(args), self.logger)
+
+        if torch.cuda.is_available():
+            self.device = torch.device("cuda")
+            self.device_count = torch.cuda.device_count()
+            assert self.device_count == 1
+            self.logger.info('using {} gpus'.format(self.device_count))
+        else:
+            raise Exception("gpu is not available")
+
+        downsample_ratio = 8
+        if args.dataset.lower() == 'qnrf':
+            self.datasets = {x: Crowd_qnrf(os.path.join(args.data_dir, x),
+                                           args.crop_size, downsample_ratio, x) for x in ['train', 'val']}
+        elif args.dataset.lower() == 'nwpu':
+            self.datasets = {x: Crowd_nwpu(os.path.join(args.data_dir, x),
+                                           args.crop_size, downsample_ratio, x) for x in ['train', 'val']}
+        elif args.dataset.lower() == 'sha' or args.dataset.lower() == 'shb':
+            self.datasets = {'train': Crowd_sh(os.path.join(args.data_dir, 'train_data'),
+                                               args.crop_size, downsample_ratio, 'train'),
+                             'val': Crowd_sh(os.path.join(args.data_dir, 'test_data'),
+                                             args.crop_size, downsample_ratio, 'val'),
+                             }
+        else:
+            raise NotImplementedError
+
+        self.dataloaders = {x: DataLoader(self.datasets[x],
+                                          collate_fn=(train_collate
+                                                      if x == 'train' else default_collate),
+                                          batch_size=(args.batch_size
+                                                      if x == 'train' else 1),
+                                          shuffle=(True if x == 'train' else False),
+                                          num_workers=args.num_workers * self.device_count,
+                                          pin_memory=(True if x == 'train' else False))
+                            for x in ['train', 'val']}
+        self.model = vgg19()
+        self.model.to(self.device)
+        self.optimizer = optim.Adam(self.model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+
+        self.start_epoch = 0
+        if args.resume:
+            self.logger.info('loading pretrained model from ' + args.resume)
+            suf = args.resume.rsplit('.', 1)[-1]
+            if suf == 'tar':
+                checkpoint = torch.load(args.resume, self.device)
+                self.model.load_state_dict(checkpoint['model_state_dict'])
+                self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+                self.start_epoch = checkpoint['epoch'] + 1
+            elif suf == 'pth':
+                self.model.load_state_dict(torch.load(args.resume, self.device))
+        else:
+            self.logger.info('random initialization')
+
+        self.ot_loss = OT_Loss(args.crop_size, downsample_ratio, args.norm_cood, self.device, args.num_of_iter_in_ot,
+                               args.reg)
+        self.tv_loss = nn.L1Loss(reduction='none').to(self.device)
+        self.mse = nn.MSELoss().to(self.device)
+        self.mae = nn.L1Loss().to(self.device)
+        self.save_list = Save_Handle(max_num=1)
+        self.best_mae = np.inf
+        self.best_mse = np.inf
+        self.best_count = 0
+
+    def train(self):
+        """training process"""
+        args = self.args
+        for epoch in range(self.start_epoch, args.max_epoch + 1):
+            self.logger.info('-' * 5 + 'Epoch {}/{}'.format(epoch, args.max_epoch) + '-' * 5)
+            self.epoch = epoch
+            self.train_eopch()
+            if epoch % args.val_epoch == 0 and epoch >= args.val_start:
+                self.val_epoch()
+
+    def train_eopch(self):
+        epoch_ot_loss = AverageMeter()
+        epoch_ot_obj_value = AverageMeter()
+        epoch_wd = AverageMeter()
+        epoch_count_loss = AverageMeter()
+        epoch_tv_loss = AverageMeter()
+        epoch_loss = AverageMeter()
+        epoch_mae = AverageMeter()
+        epoch_mse = AverageMeter()
+        epoch_start = time.time()
+        self.model.train()  # Set model to training mode
+
+        for step, (inputs, points, gt_discrete) in enumerate(self.dataloaders['train']):
+            inputs = inputs.to(self.device)
+            gd_count = np.array([len(p) for p in points], dtype=np.float32)
+            points = [p.to(self.device) for p in points]
+            gt_discrete = gt_discrete.to(self.device)
+            N = inputs.size(0)
+
+            with torch.set_grad_enabled(True):
+                outputs, outputs_normed = self.model(inputs)
+                # Compute OT loss.
+                ot_loss, wd, ot_obj_value = self.ot_loss(outputs_normed, outputs, points)
+                ot_loss = ot_loss * self.args.wot
+                ot_obj_value = ot_obj_value * self.args.wot
+                epoch_ot_loss.update(ot_loss.item(), N)
+                epoch_ot_obj_value.update(ot_obj_value.item(), N)
+                epoch_wd.update(wd, N)
+
+                # Compute counting loss.
+                count_loss = self.mae(outputs.sum(1).sum(1).sum(1),
+                                      torch.from_numpy(gd_count).float().to(self.device))
+                epoch_count_loss.update(count_loss.item(), N)
+
+                # Compute TV loss.
+                gd_count_tensor = torch.from_numpy(gd_count).float().to(self.device).unsqueeze(1).unsqueeze(
+                    2).unsqueeze(3)
+                gt_discrete_normed = gt_discrete / (gd_count_tensor + 1e-6)
+                tv_loss = (self.tv_loss(outputs_normed, gt_discrete_normed).sum(1).sum(1).sum(
+                    1) * torch.from_numpy(gd_count).float().to(self.device)).mean(0) * self.args.wtv
+                epoch_tv_loss.update(tv_loss.item(), N)
+
+                loss = ot_loss + count_loss + tv_loss
+
+                self.optimizer.zero_grad()
+                loss.backward()
+                self.optimizer.step()
+
+                pred_count = torch.sum(outputs.view(N, -1), dim=1).detach().cpu().numpy()
+                pred_err = pred_count - gd_count
+                epoch_loss.update(loss.item(), N)
+                epoch_mse.update(np.mean(pred_err * pred_err), N)
+                epoch_mae.update(np.mean(abs(pred_err)), N)
+
+        self.logger.info(
+            'Epoch {} Train, Loss: {:.2f}, OT Loss: {:.2e}, Wass Distance: {:.2f}, OT obj value: {:.2f}, '
+            'Count Loss: {:.2f}, TV Loss: {:.2f}, MSE: {:.2f} MAE: {:.2f}, Cost {:.1f} sec'
+                .format(self.epoch, epoch_loss.get_avg(), epoch_ot_loss.get_avg(), epoch_wd.get_avg(),
+                        epoch_ot_obj_value.get_avg(), epoch_count_loss.get_avg(), epoch_tv_loss.get_avg(),
+                        np.sqrt(epoch_mse.get_avg()), epoch_mae.get_avg(),
+                        time.time() - epoch_start))
+        model_state_dic = self.model.state_dict()
+        save_path = os.path.join(self.save_dir, '{}_ckpt.tar'.format(self.epoch))
+        torch.save({
+            'epoch': self.epoch,
+            'optimizer_state_dict': self.optimizer.state_dict(),
+            'model_state_dict': model_state_dic
+        }, save_path)
+        self.save_list.append(save_path)
+
+    def val_epoch(self):
+        args = self.args
+        epoch_start = time.time()
+        self.model.eval()  # Set model to evaluate mode
+        epoch_res = []
+        for inputs, count, name in self.dataloaders['val']:
+            inputs = inputs.to(self.device)
+            assert inputs.size(0) == 1, 'the batch size should equal to 1 in validation mode'
+            with torch.set_grad_enabled(False):
+                outputs, _ = self.model(inputs)
+                res = count[0].item() - torch.sum(outputs).item()
+                epoch_res.append(res)
+
+        epoch_res = np.array(epoch_res)
+        mse = np.sqrt(np.mean(np.square(epoch_res)))
+        mae = np.mean(np.abs(epoch_res))
+        self.logger.info('Epoch {} Val, MSE: {:.2f} MAE: {:.2f}, Cost {:.1f} sec'
+                         .format(self.epoch, mse, mae, time.time() - epoch_start))
+
+        model_state_dic = self.model.state_dict()
+        if (2.0 * mse + mae) < (2.0 * self.best_mse + self.best_mae):
+            self.best_mse = mse
+            self.best_mae = mae
+            self.logger.info("save best mse {:.2f} mae {:.2f} model epoch {}".format(self.best_mse,
+                                                                                     self.best_mae,
+                                                                                     self.epoch))
+            torch.save(model_state_dic, os.path.join(self.save_dir, 'best_model_{}.pth'.format(self.best_count)))
+            self.best_count += 1

code/utils/log_utils.py

@@ -0,0 +1,24 @@
+import logging
+
+
+def get_logger(log_file):
+    logger = logging.getLogger(log_file)
+    logger.setLevel(logging.DEBUG)
+    fh = logging.FileHandler(log_file)
+    fh.setLevel(logging.DEBUG)
+    ch = logging.StreamHandler()
+    ch.setLevel(logging.INFO)
+    formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
+    ch.setFormatter(formatter)
+    fh.setFormatter(formatter)
+    logger.addHandler(ch)
+    logger.addHandler(fh)
+    return logger
+
+
+def print_config(config, logger):
+    """
+    Print configuration of the model
+    """
+    for k, v in config.items():
+        logger.info("{}:\t{}".format(k.ljust(15), v))

code/utils/pytorch_utils.py

@@ -0,0 +1,58 @@
+import os
+
+def adjust_learning_rate(optimizer, epoch, initial_lr=0.001, decay_epoch=10):
+    """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
+    lr = max(initial_lr * (0.1 ** (epoch // decay_epoch)), 1e-6)
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+
+class Save_Handle(object):
+    """handle the number of """
+    def __init__(self, max_num):
+        self.save_list = []
+        self.max_num = max_num
+
+    def append(self, save_path):
+        if len(self.save_list) < self.max_num:
+            self.save_list.append(save_path)
+        else:
+            remove_path = self.save_list[0]
+            del self.save_list[0]
+            self.save_list.append(save_path)
+            if os.path.exists(remove_path):
+                os.remove(remove_path)
+
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = 1.0 * self.sum / self.count
+
+    def get_avg(self):
+        return self.avg
+
+    def get_count(self):
+        return self.count
+
+
+def set_trainable(model, requires_grad):
+    for param in model.parameters():
+        param.requires_grad = requires_grad
+
+
+
+def get_num_params(model):
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)