├── README.md ├── dataset.py ├── image.py ├── model.py ├── network.py ├── part_A_test.json ├── test.py └── utils.py /README.md: -------------------------------------------------------------------------------- 1 | 2 | # Residual Regression with Semantic Prior for Crowd Counting 3 | 4 | This is the official code of the CVPR 2019 PAPER. 5 | 6 | [CVPR 2019 PAPER](http://visal.cs.cityu.edu.hk/static/pubs/conf/cvpr19-rescount.pdf) | [Project Page](http://visal.cs.cityu.edu.hk/research/residual_regression_counting/) 7 | 8 | ## Usage 9 | 10 | 1. Requirements: 11 | * python2.7 12 | * pytorch 0.4.1 13 | 2. Testing: 14 | * download [test set](https://drive.google.com/open?id=11WJwj5EDxAgkZwJV47gmiqm3bL4hVNLr) 15 | * download [pretrained models](https://drive.google.com/drive/folders/1CBGA47LxGM1UToWcbn3oVusk0iEbM3cZ) 16 | * run `python test.py` 17 | 18 | ## Citing 19 | ``` 20 | @inproceedings{wan2019residual, 21 | title={Residual Regression With Semantic Prior for Crowd Counting}, 22 | author={Wan, Jia and Luo, Wenhan and Wu, Baoyuan and Chan, Antoni B and Liu, Wei}, 23 | booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition}, 24 | pages={4036--4045}, 25 | year={2019} 26 | } 27 | ``` 28 | 29 | -------------------------------------------------------------------------------- /dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import torch 4 | import numpy as np 5 | from torch.utils.data import Dataset 6 | from PIL import Image 7 | from image import load_data 8 | import torchvision.transforms.functional as F 9 | 10 | class listDataset(Dataset): 11 | def __init__(self, root, transform, num_workers=4): 12 | self.lines = root 13 | self.nSamples = len(root) 14 | self.transform = transform 15 | 16 | def __len__(self): 17 | return self.nSamples 18 | 19 | def __getitem__(self, index): 20 | assert index <= len(self), 'index range error' 21 | 22 | img_path = self.lines[index] 23 | 24 | img,target,smap = load_data(img_path) 25 | img = self.transform(img) 26 | 27 | return img,target,smap 28 | 29 | -------------------------------------------------------------------------------- /image.py: -------------------------------------------------------------------------------- 1 | from PIL import Image 2 | import numpy as np 3 | import h5py 4 | import cv2 5 | 6 | 7 | def load_data(img_path): 8 | 9 | img = Image.open(img_path).convert('RGB') 10 | w, h = img.size 11 | d_size = 128 12 | w = int(max(d_size, np.floor(w/d_size)*d_size)) 13 | h = int(max(d_size, np.floor(h/d_size)*d_size)) 14 | img = img.resize([w, h]) 15 | 16 | gt_path = img_path.replace('.jpg', '.h5').replace('images', 'ground-truth') 17 | gt_file = h5py.File(gt_path, 'r') 18 | target = np.asarray(gt_file['density']) 19 | 20 | mask_path = img_path.replace('.jpg', '_mask.h5').replace('images', 'ground-truth') 21 | s_file = h5py.File(mask_path, 'r') 22 | smap = np.asarray(s_file['density']) 23 | smap = cv2.resize(smap, (int(w/8), int(h/8))) 24 | 25 | return img, target, smap 26 | -------------------------------------------------------------------------------- /model.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | import numpy as np 3 | import torch 4 | import math 5 | import torch.nn.functional as F 6 | from torchvision import models 7 | from utils import save_net,load_net 8 | from network import Conv2d 9 | 10 | 11 | 12 | class Residual(nn.Module): 13 | ''' 14 | given appearance estimation, deep features, and residual maps as input, 15 | output the final estimation 16 | ''' 17 | def __init__(self, n=2, k_nn=3, bn=False): 18 | super(Residual, self).__init__() 19 | dim = 64 20 | self.r1 = nn.Sequential(Conv2d(n*64,32,3,NL='prelu',same_padding=True,bn=bn)) 21 | self.r2 = nn.Sequential(Conv2d(n*64,16,5,NL='prelu',same_padding=True,bn=bn)) 22 | self.r3 = nn.Sequential(Conv2d(n*64,8,7,NL='prelu',same_padding=True,bn=bn)) 23 | self.residual_predict = nn.Sequential( 24 | Conv2d(56, 16, 7, NL='prelu', same_padding=True, bn=bn), 25 | Conv2d(16, 8, 5, NL='prelu', same_padding=True, bn=bn), 26 | Conv2d(8, 1, 3, NL='nrelu', same_padding=True, bn=bn)) 27 | 28 | self.rm1 = nn.Sequential(Conv2d(k_nn,16,1,NL='relu',same_padding=True,bn=bn)) 29 | self.rm2 = nn.Sequential(Conv2d(k_nn,8,3,NL='relu',same_padding=True,bn=bn)) 30 | self.rm3 = nn.Sequential(Conv2d(k_nn,4,5,NL='relu',same_padding=True,bn=bn)) 31 | self.residual_merge = nn.Sequential(Conv2d(28,1,3,NL='relu',same_padding=True,bn=bn)) 32 | 33 | self.ram1 = nn.Sequential(Conv2d(2,16,1,NL='relu',same_padding=True,bn=bn)) 34 | self.ram2 = nn.Sequential(Conv2d(2,8,3,NL='relu',same_padding=True,bn=bn)) 35 | self.ram3 = nn.Sequential(Conv2d(2,4,5,NL='relu',same_padding=True,bn=bn)) 36 | self.res_app_merge = nn.Sequential(Conv2d(28,1,3,NL='relu',same_padding=True,bn=bn)) 37 | 38 | self._initialize_weights() 39 | 40 | def forward(self, features, app_prediction, support_gt): 41 | # pair test image and support images 42 | for i in range(len(features)-1): 43 | pair = torch.cat((features[0],features[i+1]),0) 44 | if i == 0: 45 | pairs = pair.unsqueeze(0) 46 | else: 47 | pairs = torch.cat((pairs, pair.unsqueeze(0)),0) 48 | # predict residual maps 49 | x1 = self.r1(pairs) 50 | x2 = self.r2(pairs) 51 | x3 = self.r3(pairs) 52 | pairs = self.residual_predict(torch.cat((x1,x2,x3),1)) 53 | # calc residual based density estimation 54 | residual_predictions = pairs + support_gt.squeeze(0).unsqueeze(1) 55 | # merge residual mals 56 | n = len(residual_predictions) 57 | h = residual_predictions.shape[2] 58 | w = residual_predictions.shape[3] 59 | x1 = self.rm1(residual_predictions.view(1,n,h,w)) 60 | x2 = self.rm2(residual_predictions.view(1,n,h,w)) 61 | x3 = self.rm3(residual_predictions.view(1,n,h,w)) 62 | final_residual_prediction = self.residual_merge(torch.cat((x1,x2,x3),1)) 63 | # merge residual and appearance maps 64 | x1 = self.ram1(torch.cat((final_residual_prediction, app_prediction),1)) 65 | x2 = self.ram2(torch.cat((final_residual_prediction, app_prediction),1)) 66 | x3 = self.ram3(torch.cat((final_residual_prediction, app_prediction),1)) 67 | final_prediction = self.res_app_merge(torch.cat((x1,x2,x3),1)) 68 | #final_prediction = self.res_app_merge(torch.cat((app_prediction,final_residual_prediction),1)) 69 | return final_prediction, residual_predictions, final_residual_prediction 70 | 71 | def _initialize_weights(self): 72 | for m in self.modules(): 73 | if isinstance(m, nn.Conv2d): 74 | nn.init.normal_(m.weight, std=0.01) 75 | if m.bias is not None: 76 | nn.init.constant_(m.bias, 0) 77 | elif isinstance(m, nn.BatchNorm2d): 78 | nn.init.constant_(m.weight, 1) 79 | nn.init.constant_(m.bias, 0) 80 | 81 | class CSRNet(nn.Module): 82 | def __init__(self, load_weights=False,norm=False,dropout=False): 83 | super(CSRNet, self).__init__() 84 | self.seen = 0 85 | self.frontend_feat = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512] 86 | self.backend_feat = [512, 512, 512,256,128,64] 87 | self.frontend = make_layers(self.frontend_feat) 88 | self.backend = make_layers(self.backend_feat,in_channels = 512,norm=norm,dilation = True, dropout=dropout) 89 | self.output_layer = nn.Conv2d(64, 1, kernel_size=1) 90 | self._initialize_weights() 91 | if not load_weights: 92 | mod = models.vgg16(pretrained = True) 93 | fs = self.frontend.state_dict() 94 | ms = mod.state_dict() 95 | for key in fs: 96 | fs[key] = ms['features.'+key] 97 | self.frontend.load_state_dict(fs) 98 | else: 99 | print("Don't pre-train on ImageNet") 100 | 101 | def forward(self,x): 102 | x = self.frontend(x) 103 | x = self.backend(x) 104 | self.features = x 105 | x = self.output_layer(x) 106 | #x = F.interpolate(x, scale_factor=8) 107 | return x 108 | def _initialize_weights(self): 109 | for m in self.modules(): 110 | if isinstance(m, nn.Conv2d): 111 | nn.init.normal_(m.weight, std=0.01) 112 | if m.bias is not None: 113 | nn.init.constant_(m.bias, 0) 114 | elif isinstance(m, nn.BatchNorm2d): 115 | nn.init.constant_(m.weight, 1) 116 | nn.init.constant_(m.bias, 0) 117 | elif isinstance(m, nn.ConvTranspose2d): 118 | nn.init.normal_(m.weight, std=0.01) 119 | if m.bias is not None: 120 | nn.init.constant_(m.bias, 0) 121 | 122 | 123 | def make_layers(cfg, in_channels = 3,norm=False,dilation = False, dropout=False): 124 | if dilation: 125 | d_rate = 2 126 | else: 127 | d_rate = 1 128 | layers = [] 129 | for v in cfg: 130 | if v == 'M': 131 | layers += [nn.MaxPool2d(kernel_size=2, stride=2)] 132 | else: 133 | conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=d_rate,dilation = d_rate) 134 | if norm: 135 | layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)] 136 | else: 137 | layers += [conv2d, nn.ReLU(inplace=True)] 138 | in_channels = v 139 | return nn.Sequential(*layers) 140 | -------------------------------------------------------------------------------- /network.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.autograd import Variable 4 | import numpy as np 5 | import cv2 6 | 7 | class Conv2d(nn.Module): 8 | def __init__(self, in_channels, out_channels, kernel_size, stride=1, NL='relu', same_padding=False, bn=False): 9 | super(Conv2d, self).__init__() 10 | padding = int((kernel_size - 1) / 2) if same_padding else 0 11 | self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding=padding) 12 | self.bn = nn.BatchNorm2d(out_channels, eps=0.001, momentum=0.1, affine=True) if bn else None 13 | if NL == 'relu': 14 | self.relu = nn.ReLU(inplace=True) 15 | elif NL == 'prelu': 16 | self.relu = nn.PReLU() 17 | elif NL == 'lrelu': 18 | self.relu = nn.LeakyReLU(0.2, True) 19 | else: 20 | self.relu = None 21 | 22 | def forward(self, x): 23 | x = self.conv(x) 24 | if self.bn is not None: 25 | x = self.bn(x) 26 | if self.relu is not None: 27 | x = self.relu(x) 28 | return x 29 | 30 | 31 | class FC(nn.Module): 32 | def __init__(self, in_features, out_features, relu=True): 33 | super(FC, self).__init__() 34 | self.fc = nn.Linear(in_features, out_features) 35 | self.relu = nn.ReLU(inplace=True) if relu else None 36 | 37 | def forward(self, x): 38 | x = self.fc(x) 39 | if self.relu is not None: 40 | x = self.relu(x) 41 | return x 42 | 43 | 44 | def save_net(fname, net): 45 | import h5py 46 | with h5py.File(fname, mode='w') as h5f: 47 | for k, v in net.state_dict().items(): 48 | #if 'module' in k: 49 | # k = k.replace('.module','.') 50 | h5f.create_dataset(k, data=v.cpu().numpy()) 51 | 52 | 53 | def load_net(fname, net): 54 | import h5py 55 | with h5py.File(fname, mode='r') as h5f: 56 | for k, v in net.state_dict().items(): 57 | #print(k) 58 | param = torch.from_numpy(np.asarray(h5f[k])) 59 | v.copy_(param) 60 | 61 | 62 | def np_to_variable(x, is_cuda=True, is_training=False, dtype=torch.FloatTensor): 63 | if is_training: 64 | v = Variable(torch.from_numpy(x).type(dtype)) 65 | else: 66 | v = Variable(torch.from_numpy(x).type(dtype), requires_grad = False, volatile = True) 67 | if is_cuda: 68 | v = v.cuda() 69 | return v 70 | 71 | def to_variable(x, is_cuda=True, is_training=False): 72 | if is_training: 73 | v = Variable(x) 74 | else: 75 | v = Variable(x, requires_grad = False, volatile = True) 76 | if is_cuda: 77 | v = v.cuda() 78 | return v 79 | 80 | 81 | def set_trainable(model, requires_grad): 82 | for param in model.parameters(): 83 | param.requires_grad = requires_grad 84 | 85 | 86 | def weights_normal_init(model, dev=0.01): 87 | if isinstance(model, list): 88 | for m in model: 89 | weights_normal_init(m, dev) 90 | else: 91 | for m in model.modules(): 92 | if isinstance(m, nn.Conv2d): 93 | #print torch.sum(m.weight) 94 | m.weight.data.normal_(0.0, dev) 95 | if m.bias is not None: 96 | m.bias.data.fill_(0.0) 97 | elif isinstance(m, nn.Linear): 98 | m.weight.data.normal_(0.0, dev) 99 | elif isinstance(m, nn.PReLU): 100 | m.weight.data.fill_(0.001) 101 | 102 | 103 | # put i-th map to the top 104 | def to_top(data, i): 105 | temp = data 106 | if(i == 0): 107 | return temp 108 | temp[0,:,:,:,] = data[i,:,:,:] 109 | temp[i,:,:,:,] = data[0,:,:,:] 110 | return temp 111 | 112 | -------------------------------------------------------------------------------- /part_A_test.json: -------------------------------------------------------------------------------- 1 | ["./datasets/ShanghaiTechA/test_data/images/IMG_73.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_38.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_54.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_145.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_70.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_91.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_143.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_173.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_100.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_3.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_95.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_104.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_18.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_181.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_151.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_169.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_121.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_14.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_129.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_174.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_109.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_172.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_71.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_74.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_161.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_43.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_28.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_67.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_152.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_130.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_93.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_156.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_163.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_12.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_177.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_46.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_120.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_24.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_47.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_9.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_167.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_141.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_136.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_15.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_148.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_112.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_8.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_164.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_65.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_44.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_61.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_179.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_56.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_131.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_85.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_139.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_89.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_26.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_128.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_32.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_119.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_60.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_165.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_166.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_154.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_64.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_59.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_78.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_135.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_114.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_4.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_175.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_96.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_58.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_146.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_92.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_83.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_105.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_34.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_111.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_101.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_17.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_22.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_107.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_103.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_127.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_106.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_33.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_125.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_144.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_138.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_29.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_155.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_115.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_149.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_72.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_134.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_153.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_133.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_116.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_51.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_123.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_41.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_81.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_55.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_19.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_48.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_7.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_158.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_80.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_118.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_84.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_76.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_90.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_30.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_53.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_176.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_68.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_168.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_63.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_86.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_20.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_147.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_108.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_45.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_82.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_10.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_49.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_5.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_87.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_160.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_170.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_178.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_126.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_162.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_23.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_2.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_1.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_171.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_52.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_42.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_21.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_132.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_31.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_6.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_140.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_40.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_142.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_36.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_110.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_77.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_27.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_137.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_57.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_159.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_122.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_113.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_62.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_124.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_150.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_39.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_99.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_37.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_75.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_157.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_98.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_66.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_94.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_13.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_102.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_180.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_11.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_25.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_97.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_182.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_35.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_50.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_88.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_79.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_69.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_117.jpg", "./datasets/ShanghaiTechA/test_data/images/IMG_16.jpg"] 2 | -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | # This is the testing script for "Residual Regression and Semantic Prior for crowd counting" 2 | from model import Residual, CSRNet 3 | import torch 4 | import torch.nn as nn 5 | from torchvision import datasets, transforms 6 | 7 | import numpy as np 8 | import argparse 9 | import json 10 | import cv2 11 | import dataset 12 | 13 | parser = argparse.ArgumentParser(description='Residual regression') 14 | parser.add_argument('--test_json', metavar='TEST', type=str, default='part_A_test.json', 15 | help='path to test json') 16 | parser.add_argument('--n', metavar='Number', type=int, default=3, 17 | help='number of support images') 18 | parser.add_argument('--thr', metavar='THRESHOLD', type=float, default=0.7, 19 | help='threshold for semantic prior') 20 | 21 | transform=transforms.Compose([transforms.ToTensor(), 22 | transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) 23 | 24 | def main(): 25 | 26 | global args 27 | global support_imgs, support_gt, support_imgs_fea 28 | 29 | args = parser.parse_args() 30 | with open(args.test_json, 'r') as outfile: 31 | val_list = json.load(outfile) 32 | 33 | residual = Residual(k_nn=args.n) 34 | residual = residual.cuda() 35 | 36 | counter = CSRNet() 37 | counter = counter.cuda() 38 | 39 | # load counter params 40 | checkpoint = torch.load('./saved_models/countercheckpoint.pth.tar') 41 | counter.load_state_dict(checkpoint['state_dict_model']) 42 | 43 | # load residual regressor params 44 | checkpoint = torch.load('./saved_models/residualcheckpoint.pth.tar') 45 | residual.load_state_dict(checkpoint['state_dict_res']) 46 | support_imgs = checkpoint['support_imgs'].cuda() 47 | support_gt = checkpoint['support_gt'].cuda() 48 | 49 | counter(support_imgs) 50 | support_imgs_fea = counter.features 51 | 52 | app, res, final = validate(val_list, counter, residual) 53 | 54 | 55 | def validate(val_list, counter, residual): 56 | print ('begin test') 57 | test_loader = torch.utils.data.DataLoader( 58 | dataset.listDataset(val_list, 59 | transform=transform), 60 | batch_size=1) 61 | 62 | app_mae = 0 63 | res_mae = 0 64 | final_mae = 0 65 | 66 | i = -1 67 | for data in (test_loader): 68 | i += 1 69 | img, target, smap = data 70 | img = img.cuda() 71 | output = counter(img) 72 | output_fea = counter.features 73 | 74 | n_h = int(output_fea.shape[2]/support_imgs_fea.shape[2]) 75 | n_w = int(output_fea.shape[3]/support_imgs_fea.shape[3]) 76 | fea = torch.cat((output_fea, support_imgs_fea.repeat(1,1,n_h,n_w)),0) 77 | # CSRNet is used to extract image features 78 | final_output, _, final_residual = residual(fea, output, support_gt.repeat(1,1,n_h,n_w)) 79 | 80 | # semantic map 81 | smap = smap.type(torch.cuda.FloatTensor).unsqueeze(0) 82 | smap[smap > 0] = 1 83 | smap[smap <= 0] = args.thr 84 | 85 | final_output = torch.mul(final_output, smap) 86 | 87 | app_mae += abs(output[0].data.sum()-target.sum().type(torch.FloatTensor).cuda()) 88 | res_mae += abs(final_residual.data.sum()-target.sum().type(torch.FloatTensor).cuda()) 89 | final_mae += abs(final_output.data.sum()-target.sum().type(torch.FloatTensor).cuda()) 90 | 91 | app_mae = app_mae/len(test_loader) 92 | res_mae = res_mae/len(test_loader) 93 | final_mae = final_mae/len(test_loader) 94 | print(' * app MAE {mae:.3f} * residual MAE {res_mae:.3f} * final MAE {final_mae:.3f}' 95 | .format(mae=app_mae, res_mae=res_mae, final_mae=final_mae)) 96 | 97 | return app_mae, res_mae, final_mae 98 | 99 | 100 | if __name__ == '__main__': 101 | main() 102 | -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import h5py 2 | import torch 3 | import shutil 4 | 5 | def save_net(fname, net): 6 | with h5py.File(fname, 'w') as h5f: 7 | for k, v in net.state_dict().items(): 8 | h5f.create_dataset(k, data=v.cpu().numpy()) 9 | def load_net(fname, net): 10 | with h5py.File('./saved_models'+fname, 'r') as h5f: 11 | for k, v in net.state_dict().items(): 12 | param = torch.from_numpy(np.asarray(h5f[k])) 13 | v.copy_(param) 14 | 15 | def save_checkpoint(state, is_best,task_id, filename='checkpoint.pth.tar'): 16 | torch.save(state, './saved_models/'+task_id+filename) 17 | if is_best: 18 | shutil.copyfile('./saved_models/'+task_id+filename, './saved_models/'+task_id+'model_best.pth.tar') 19 | 20 | def tv_loss(y): 21 | loss = torch.sum((y[:-1, :, :, :] - y[1:, :, :, :])**2) + torch.sum((y[:, :-1, :, :] - y[:, 1:, :, :])**2) 22 | return loss 23 | 24 | --------------------------------------------------------------------------------