class CC:
# class for crowd counting with SANet
def __init__(self, args):
# setup the net
torch.device("cuda")
torch.cuda.set_device(0)
torch.backends.cudnn.enabled = True
self.net = CrowdCounter(cfg.GPU_ID, cfg.NET, torch.nn.MSELoss(),
pytorch_ssim.SSIM(window_size=11))
self.net.load_state_dict(torch.load(args.model))
self.net.cuda()
self.net.eval()
# IMPORTANT: when changing model, make sure you change config for the dataset it was trained in to make sure mean, std, h and w are correct
self.mean, self.std = cfg.MEAN_STD
self.h, self.w = cfg.STD_SIZE
print(self.net) # print net structure
# video capture object
self.vid = VideoCapture(args.video)
def run(self, thread_queue):
# thread function to run
global go_thread
while go_thread:
ret, frame = self.vid.get_frame()
if ret:
# resize, convert to pytorch tensor, normalize
if frame.shape[0] != self.h or frame.shape[1] != self.w:
frame = cv2.resize(frame, (self.w, self.h),
interpolation=cv2.INTER_CUBIC)
tensor = torchvision.transforms.ToTensor()(frame)
tensor = torchvision.transforms.functional.normalize(
tensor, mean=self.mean, std=self.std)
# forward propagation
with torch.no_grad():
tensor = torch.autograd.Variable(
tensor[None, :, :, :]).cuda()
pred_map = self.net.test_forward(tensor)
pred_map = pred_map.cpu().data.numpy()[0, 0, :, :]
# generate grayscale image for overlap and put results in thread queue
gray = np.repeat(cv2.cvtColor(frame,
cv2.COLOR_RGB2GRAY)[:, :,
np.newaxis],
3,
axis=2)
thread_queue.put((gray, pred_map))
class CC:
""" Class for crowd counting with SANet (threaded) """
def __init__(self, model):
# Setup the net
torch.device("cuda")
torch.cuda.set_device(0)
torch.backends.cudnn.enabled = True
self.net = CrowdCounter(cfg.GPU_ID, cfg.NET, torch.nn.MSELoss(),
pytorch_ssim.SSIM(window_size=11))
self.net.load_state_dict(torch.load(model))
self.net.cuda()
self.net.eval()
# IMPORTANT: when changing model, make sure you change config for the dataset it was trained in to make sure mean and std are correct
self.mean, self.std = cfg.MEAN_STD
# Print net structure
# print(self.net)
def run(self, stop_event, cam_queue, net_queue):
while not stop_event.is_set():
try:
frame = cam_queue.get(0)
# convert to pytorch tensor, normalize
tensor = torchvision.transforms.ToTensor()(frame)
tensor = torchvision.transforms.functional.normalize(
tensor, mean=self.mean, std=self.std)
# padding to multiples of 8
xr = (8 - tensor.shape[2] % 8) % 8
yr = (8 - tensor.shape[1] % 8) % 8
tensor = torch.nn.functional.pad(tensor, (xr, xr, yr, yr),
'constant', 0)
# forward propagation
with torch.no_grad():
tensor = torch.autograd.Variable(
tensor[None, :, :, :]).cuda()
pred_map = self.net.test_forward(tensor)
pred_map = pred_map.cpu().data.numpy()[0, 0, :, :]
gray = np.repeat(np.array(frame.convert("L"))[:, :,
np.newaxis],
3,
axis=2)
gray = gray.astype(np.float32) / 255
net_queue.put((gray, pred_map))
except queue.Empty:
pass
# open the video stream / file
cap = cv2.VideoCapture(args.video)
while (cap.isOpened()):
_, frame = cap.read()
if frame is None:
break
# convert to pytorch tensor and normalize
tensor = torchvision.transforms.ToTensor()(cv2.cvtColor(
frame, cv2.COLOR_BGR2RGB))
tensor = torchvision.transforms.functional.normalize(tensor,
mean=mean,
std=std)
# forward propagation
with torch.no_grad():
tensor = torch.autograd.Variable(tensor[None, :, :, :]).cuda()
pred_map = net.test_forward(tensor)
pred_map = pred_map.cpu().data.numpy()[0, 0, :, :]
# converts frame to grayscale and density map to color map, then overlap them
gray = np.repeat(cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)[:, :,
np.newaxis],
3,
axis=2)
dmap = cv2.applyColorMap(
np.array(pred_map / np.max(pred_map + 1e-20) * 255,
dtype=np.uint8), args.cmap)
dmap = dmap.astype(np.float32) / 255
gray = gray.astype(np.float32) / 255
if args.mode == 'Add':
res = dmap * args.opacity + gray
elif args.mode == 'Lighten':
res = np.maximum(gray, dmap * args.opacity)
def test(file_list, model_path):
loss_1_fn = torch.nn.MSELoss()
loss_2_fn = pytorch_ssim.SSIM(window_size=11)
net = CrowdCounter(cfg.GPU_ID, cfg.NET, loss_1_fn, loss_2_fn)
net.load_state_dict(torch.load(model_path))
net.cuda()
net.eval()
f1 = plt.figure(1)
gts = []
preds = []
for filename in file_list:
print(filename, end=', ')
imgname = dataRoot + '/img/' + filename
filename_no_ext = filename.split('.')[0]
denname = dataRoot + '/den/' + filename_no_ext + '.csv'
den = pd.read_csv(denname, sep=',', header=None).values
den = den.astype(np.float32, copy=False)
img = Image.open(imgname)
if img.mode == 'L':
img = img.convert('RGB')
img = img_transform(img)
if slicing:
xr = (8 - img.shape[2] % 8) % 8
yr = (8 - img.shape[1] % 8) % 8
img = torch.nn.functional.pad(img, (xr, xr, yr, yr), 'constant', 0)
pred_maps = []
x4 = img.shape[2] # full image
x1 = x4 // 2 # half image
x2 = x1 // 2 # quarter image
x3 = x1 + x2
y4 = img.shape[1]
y1 = y4 // 2
y2 = y1 // 2
y3 = y1 + y2
img_list = [
img[:, 0:y1, 0:x1], img[:, 0:y1, x2:x3], img[:, 0:y1, x1:x4],
img[:, y2:y3, 0:x1], img[:, y2:y3, x2:x3], img[:, y2:y3,
x1:x4],
img[:, y1:y4, 0:x1], img[:, y1:y4, x2:x3], img[:, y1:y4, x1:x4]
]
for inputs in img_list:
with torch.no_grad():
img = torch.autograd.Variable(inputs[None, :, :, :]).cuda()
pred_maps.append(net.test_forward(img))
x3, x5 = int(x4 * 3 / 8), int(x4 * 5 / 8)
y3, y5 = int(y4 * 3 / 8), int(y4 * 5 / 8)
x32, x52, x51, x41 = x3 - x2, x5 - x2, x5 - x1, x4 - x1
y32, y52, y51, y41 = y3 - y2, y5 - y2, y5 - y1, y4 - y1
slice0 = pred_maps[0].cpu().data.numpy()[0, 0, 0:y3, 0:x3]
slice1 = pred_maps[1].cpu().data.numpy()[0, 0, 0:y3, x32:x52]
slice2 = pred_maps[2].cpu().data.numpy()[0, 0, 0:y3, x51:x41]
slice3 = pred_maps[3].cpu().data.numpy()[0, 0, y32:y52, 0:x3]
slice4 = pred_maps[4].cpu().data.numpy()[0, 0, y32:y52, x32:x52]
slice5 = pred_maps[5].cpu().data.numpy()[0, 0, y32:y52, x51:x41]
slice6 = pred_maps[6].cpu().data.numpy()[0, 0, y51:y41, 0:x3]
slice7 = pred_maps[7].cpu().data.numpy()[0, 0, y51:y41, x32:x52]
slice8 = pred_maps[8].cpu().data.numpy()[0, 0, y51:y41, x51:x41]
pred_map = np.vstack((np.hstack(
(slice0, slice1, slice2)), np.hstack((slice3, slice4, slice5)),
np.hstack((slice6, slice7, slice8))))
sio.savemat(exp_name + '/pred/' + filename_no_ext + '.mat',
{'data': pred_map / 100.})
else:
with torch.no_grad():
img = torch.autograd.Variable(img[None, :, :, :]).cuda()
pred_map = net.test_forward(img)
sio.savemat(exp_name + '/pred/' + filename_no_ext + '.mat',
{'data': pred_map.squeeze().cpu().numpy() / 100.})
pred_map = pred_map.cpu().data.numpy()[0, 0, :, :]
pred = np.sum(pred_map) / 100.0
preds.append(pred)
gt = np.sum(den)
gts.append(gt)
sio.savemat(exp_name + '/gt/' + filename_no_ext + '.mat',
{'data': den})
pred_map = pred_map / np.max(pred_map + 1e-20)
den = den / np.max(den + 1e-20)
if save_graphs:
den_frame = plt.gca()
plt.imshow(den, 'jet')
den_frame.axes.get_yaxis().set_visible(False)
den_frame.axes.get_xaxis().set_visible(False)
den_frame.spines['top'].set_visible(False)
den_frame.spines['bottom'].set_visible(False)
den_frame.spines['left'].set_visible(False)
den_frame.spines['right'].set_visible(False)
plt.savefig(exp_name + '/' + filename_no_ext + '_gt_' +
str(int(gt)) + '.png',
bbox_inches='tight',
pad_inches=0,
dpi=150)
plt.close()
# sio.savemat(exp_name+'/'+filename_no_ext+'_gt_'+str(int(gt))+'.mat',{'data':den})
pred_frame = plt.gca()
plt.imshow(pred_map, 'jet')
pred_frame.axes.get_yaxis().set_visible(False)
pred_frame.axes.get_xaxis().set_visible(False)
pred_frame.spines['top'].set_visible(False)
pred_frame.spines['bottom'].set_visible(False)
pred_frame.spines['left'].set_visible(False)
pred_frame.spines['right'].set_visible(False)
plt.savefig(exp_name + '/' + filename_no_ext + '_pred_' +
str(float(pred)) + '.png',
bbox_inches='tight',
pad_inches=0,
dpi=150)
plt.close()
# sio.savemat(exp_name+'/'+filename_no_ext+'_pred_'+str(float(pred))+'.mat',{'data':pred_map})
diff = den - pred_map
diff_frame = plt.gca()
plt.imshow(diff, 'jet')
plt.colorbar()
diff_frame.axes.get_yaxis().set_visible(False)
diff_frame.axes.get_xaxis().set_visible(False)
diff_frame.spines['top'].set_visible(False)
diff_frame.spines['bottom'].set_visible(False)
diff_frame.spines['left'].set_visible(False)
diff_frame.spines['right'].set_visible(False)
plt.savefig(exp_name + '/' + filename_no_ext + '_diff.png',
bbox_inches='tight',
pad_inches=0,
dpi=150)
plt.close()
# sio.savemat(exp_name+'/diff/'+filename_no_ext+'_diff.mat',{'data':diff})
preds = np.asarray(preds)
gts = np.asarray(gts)
print('\nMAE= ' + str(np.mean(np.abs(gts - preds))))
print('MSE= ' + str(np.sqrt(np.mean((gts - preds)**2))))