def cal_mae(img_root,gt_dmap_root,model_param_path):
'''
Calculate the MAE of the test data.
img_root: the root of test image data.
gt_dmap_root: the root of test ground truth density-map data.
model_param_path: the path of specific mcnn parameters.
'''
device=torch.device("cpu")
model=CANNet()
model.load_state_dict(torch.load(model_param_path))
model.to(device)
dataset=CrowdDataset(img_root,gt_dmap_root,8,phase='test')
dataloader=torch.utils.data.DataLoader(dataset,batch_size=1,shuffle=False)
model.eval()
mae=0
with torch.no_grad():
for i,(img,gt_dmap) in enumerate(tqdm(dataloader)):
img=img.to(device)
gt_dmap=gt_dmap.to(device)
# forward propagation
et_dmap=model(img)
mae+=abs(et_dmap.data.sum()-gt_dmap.data.sum()).item()
del img,gt_dmap,et_dmap
print("model_param_path:"+model_param_path+" mae:"+str(mae/len(dataloader)))
def estimate_density_map(img_root,gt_dmap_root,model_param_path,index):
'''
Show one estimated density-map.
img_root: the root of test image data.
gt_dmap_root: the root of test ground truth density-map data.
model_param_path: the path of specific mcnn parameters.
index: the order of the test image in test dataset.
'''
device=torch.device("cuda")
model=CANNet().to(device)
model.load_state_dict(torch.load(model_param_path))
dataset=CrowdDataset(img_root,gt_dmap_root,8,phase='test')
dataloader=torch.utils.data.DataLoader(dataset,batch_size=1,shuffle=False)
model.eval()
for i,(img,gt_dmap) in enumerate(dataloader):
if i==index:
img=img.to(device)
gt_dmap=gt_dmap.to(device)
# forward propagation
et_dmap=model(img).detach()
et_dmap=et_dmap.squeeze(0).squeeze(0).cpu().numpy()
print(et_dmap.shape)
plt.imshow(et_dmap,cmap=CM.jet)
plt.show()
break
def main_live(args):
capture = cv2.VideoCapture(args.stream)
num_gpus = torch.cuda.device_count()
device = 'cuda' if num_gpus >= 1 else 'cpu'
model = CANNet().to(device)
model.load_state_dict(
torch.load(args.weights, map_location=torch.device(device)))
transform = transforms.Compose([transforms.ToTensor()])
model.eval()
while (capture.isOpened()):
ret, frame = capture.read()
if (ret):
key = cv2.waitKey(1) & 0xFF
copy_frame = copy.copy(frame)
copy_frame = transform(copy_frame)
# putting the batch dimension on tensor
copy_frame = copy_frame.unsqueeze(0)
copy_frame = copy_frame.to(device)
result = model(copy_frame)
crowd_count = int(result.data.sum().item())
font = cv2.FONT_HERSHEY_SIMPLEX
cv2.putText(frame, "Crowd Couting: " + str(crowd_count), (10, 35),
font, 1.3, (0, 0, 255), 3, cv2.LINE_AA)
cv2.imshow("Frame", frame)
print(crowd_count)
if key == ord("q"):
break
def estimate_density_map(img_root,gt_dmap_root,model_param_path):
'''
Show one estimated density-map.
img_root: the root of test image data.
gt_dmap_root: the root of test ground truth density-map data.
model_param_path: the path of specific mcnn parameters.
index: the order of the test image in test dataset.
'''
model=CANNet().cuda()
model.load_state_dict(torch.load(model_param_path))
model.eval()
dataset=CrowdDataset(img_root,gt_dmap_root,8,phase='test')
dataloader=torch.utils.data.DataLoader(dataset,batch_size=1,shuffle=False)
print('dataloader = ',dataloader)
for img,gt_dmap,img_name in dataloader:
print('img_shape = ',img.shape)
st = time.time()
img=img.cuda()
gt_dmap=gt_dmap.cuda()
# forward propagation
et_dmap=model(img).detach()
et_dmap=et_dmap.squeeze(0).squeeze(0).cpu().numpy()
pred_frame = plt.gca()
plt.imshow(et_dmap,cmap=CM.jet)
plt.show()
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( 'result/'+str(img_name)+'_out' + '.png',bbox_inches='tight', pad_inches=0, dpi=200)
plt.close()
print('tt = ',time.time()-st)
break
parser.add_argument("--opacity",
default=0.7,
type=float,
help="Opacity value for the density map overlap")
args = parser.parse_args()
# setup the model
if args.device == "cpu":
device = torch.device("cpu")
elif args.device == "gpu":
device = torch.device("cuda")
torch.backends.cudnn.enabled = True # use cudnn?
else:
raise Exception("Unknown device. Use \"cpu\" or \"gpu\".")
model = CANNet().to(device)
model.load_state_dict(torch.load(args.model))
model.eval()
torch.no_grad()
# open the video stream / file
cap = cv2.VideoCapture(args.video)
while (cap.isOpened()):
_, frame = cap.read()
# convert to pytorch tensor and normalize
tensor = torchvision.transforms.ToTensor()(cv2.cvtColor(
frame, cv2.COLOR_BGR2RGB))
# ---------- to change after new norm training ------------------------
tensor = torchvision.transforms.functional.normalize(
tensor, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
# ---------------------------------------------------------------------
@file: demo_video.py
@time: 6/27/19 2:37 PM
@desc: test videos for serials frame
'''
import torch
import cv2
import random
import matplotlib.pyplot as plt
import matplotlib.cm as CM
from torchvision import transforms
from cannet import CANNet
from my_dataset import CrowdDataset
model = CANNet()
model_param_path = './checkpoints/cvpr2019CAN_353model.pth'
model.load_state_dict(torch.load(model_param_path))
model.cuda()
model.eval()
torch.backends.cudnn.enabled = False
def read_img(img, gt_downsample):
img = img / 255
if len(img.shape) == 2:
img = img[:, :, np.newaxis]
img = np.concatenate((img, img, img), 2)
if gt_downsample > 1:
ds_rows = int(img.shape[0] // gt_downsample)
ds_cols = int(img.shape[1] // gt_downsample)
img = cv2.resize(img,