def cal_mae(img_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("cuda")
model = CSRNet()
model.load_state_dict(torch.load(model_param_path))
model.to(device)
dataset = create_test_dataloader(img_root)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=False)
model.eval()
mae = 0
with torch.no_grad():
for i, data in enumerate(tqdm(dataloader)):
image = data['image'].cuda()
gt_densitymap = data['densitymap'].cuda()
# forward propagation
et_dmap = model(image)
mae += abs(et_dmap.data.sum() - gt_densitymap.data.sum()).item()
del image, gt_densitymap, et_dmap
print("model_param_path:" + model_param_path + " mae:" +
str(mae / len(dataloader)))
def estimate_density_map(img_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 = CSRNet().to(device)
model.load_state_dict(torch.load(model_param_path))
dataset = CrowdDataset(img_root)
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)
break
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.
'''
model = CSRNet()
model.load_state_dict(torch.load(model_param_path,
map_location=cfg.device))
model.to(cfg.device)
test_dataloader = create_test_dataloader(cfg.dataset_root) # dataloader
model.eval()
sum_mae = 0
with torch.no_grad():
for i, data in enumerate(tqdm(test_dataloader)):
image = data['image'].to(cfg.device)
gt_densitymap = data['densitymap'].to(cfg.device)
# forward propagation
et_densitymap = model(image).detach()
mae = abs(et_densitymap.data.sum() - gt_densitymap.data.sum())
sum_mae += mae.item()
# clear mem
del i, data, image, gt_densitymap, et_densitymap
torch.cuda.empty_cache()
print("model_param_path:" + model_param_path + " mae:" +
str(sum_mae / len(test_dataloader)))
def estimate_density_map(img_root, gt_dmap_root, model_param_path, index):
'''
@Mushy
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("cpu")
model = CSRNet().to(device)
model.load_state_dict(torch.load(model_param_path)) # GPU
#torch.load(model_param_path, map_location=lambda storage, loc: storage) # CPU
cfg = Config()
dataloader = create_test_dataloader(cfg.dataset_root)
model.eval()
for i, data in enumerate(dataloader):
if i == index:
img = data['image'].to(device)
# gt_dmap=gt_dmap.to(device)
gt_dmap = data['densitymap'].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.gray)
plt.show()
break
def estimate_density_map_no_gt(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.
'''
image_export_folder = 'export_images_extra'
model = CSRNet()
model.load_state_dict(torch.load(model_param_path,
map_location=cfg.device))
model.to(cfg.device)
test_dataloader = create_test_extra_dataloader(
cfg.dataset_root) # dataloader
model.eval()
with torch.no_grad():
for i, data in enumerate(tqdm(test_dataloader)):
image = data['image'].to(cfg.device)
# gt_densitymap = data['densitymap'].to(cfg.device)
# forward propagation
et_densitymap = model(image).detach()
pred_count = et_densitymap.data.sum().cpu()
# actual_count = gt_densitymap.data.sum().cpu()
actual_count = 999
et_densitymap = et_densitymap.squeeze(0).squeeze(0).cpu().numpy()
# gt_densitymap = gt_densitymap.squeeze(0).squeeze(0).cpu().numpy()
image = image[0].cpu() # denormalize(image[0].cpu())
print(et_densitymap.shape)
# et is the estimated density
plt.imshow(et_densitymap, cmap=CM.jet)
plt.savefig("{}/{}_{}_{}_{}".format(image_export_folder,
str(i).zfill(3),
str(int(pred_count)),
str(int(actual_count)),
'etdm.png'))
# # gt is the ground truth density
# plt.imshow(gt_densitymap, cmap=CM.jet)
# plt.savefig("{}/{}_{}_{}_{}".format(image_export_folder,
# str(i).zfill(3),
# str(int(pred_count)),
# str(int(actual_count)), 'gtdm.png'))
# image
plt.imshow(image.permute(1, 2, 0))
plt.savefig("{}/{}_{}_{}_{}".format(image_export_folder,
str(i).zfill(3),
str(int(pred_count)),
str(int(actual_count)),
'image.png'))
# clear mem
del i, data, image, et_densitymap, pred_count, actual_count
torch.cuda.empty_cache()
def load_model(weight):
# load model with trained weights
print('Loading model.....')
model = CSRNet()
model = model.cuda()
checkpoint = torch.load(weight)
model.load_state_dict(checkpoint['state_dict'])
model.eval()
print('Loaded.')
return model
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.
'''
cfg = Config()
device = cfg.device
model = CSRNet()
model.load_state_dict(torch.load(model_param_path)) # GPU
#torch.load(model_param_path, map_location=lambda storage, loc: storage) # CPU
model.to(device)
"""
@Mushy
Changed data loader to give path From config device
"""
dataloader = create_test_dataloader(cfg.dataset_root)
#dataloader=torch.utils.data.DataLoader(cfg.dataset_root,batch_size=1,shuffle=False)
model.eval()
mae = 0
with torch.no_grad():
for i, data in enumerate(tqdm(dataloader)):
"""
@Mushy
Changed how to access the data .
"""
img = data['image'].to(device)
#gt_dmap=gt_dmap.to(device)
gt_dmap = data['densitymap'].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)))
import numpy as np
from PIL import Image
import matplotlib as plt
from matplotlib import cm as CM
from model import CSRNet
import torch
from torchvision import transforms
import random
model = CSRNet()
checkpoint = torch.load('./model.pt', map_location='cpu')
model.load_state_dict(checkpoint['state_dict'])
model.eval()
transform = transforms.Compose([transforms.ToTensor()])
def get_prediction(file):
img = Image.open(file).convert('RGB') #Get prediction
img = transform(img)
img = img.unsqueeze(0)
output = model(img)
prediction = int(output.detach().cpu().sum().numpy())
temp = np.asarray(output.detach().cpu().reshape(
output.detach().cpu().shape[2],
output.detach().cpu().shape[3]))
loss.backward()
optimizer.step()
running_loss += loss.data.item()
epoch_loss = running_loss / totalnum
epoch_mae = running_mae / totalnum
epoch_mse = np.sqrt(running_mse / totalnum)
print('Train Loss: {:.4f} density MAE: {:.4f} density MSE: {:.4f}'.
format(epoch_loss, epoch_mae, epoch_mse))
logger.scalar_summary('Train_loss', epoch_loss, epoch)
logger.scalar_summary('Train_dens_MAE', epoch_mae, epoch)
logger.scalar_summary('Train_dens_MSE', epoch_mse, epoch)
# 验证阶段
net.eval()
running_loss = 0.0
running_mse = 0.0
running_mae = 0.0
totalnum = 0
for idx, (image, densityMap) in enumerate(val_loader):
image = image.to(device)
densityMap = densityMap.to(device)
optimizer.zero_grad()
duration = time.time()
predDensityMap = net(image)
outputs_np = predDensityMap.data.cpu().numpy()
densityMap_np = densityMap.data.cpu().numpy()
def count(path):
"""
evaluates the number of larva present in input.
input is either an image of a video. if input is an image, the evaluation is done once over the image, if input is a
video, the evaluation is done over every caption in the video seperately and then averaged over all captions to
produce the result
:param path: a path to an image or a video
:return: count
"""
# Define the device(processor) type
if torch.cuda.is_available():
device = torch.device('cuda')
print('Current procesor is GPU')
else:
device = torch.device('cpu')
print('Current procesor is CPU')
# Define the model to use for calculations
model = CSRNet()
model.load_state_dict(torch.load('model_wgts.pth'))
model.to(device)
model.eval()
# Load the image or video
im_list = []
try:
im_list.append(Image.open(path))
except OSError:
if 'http' in path:
wget.download(path, out='videos')
cap = cv2.VideoCapture(os.listdir('videos')[0])
frameCount = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
fc = 0
ret = True
im_list = []
while (fc < frameCount and ret):
ret, im = cap.read()
if fc % 10 == 0:
new_im = np.zeros_like(im)
new_im[:, :, 0] = im[:, :, 2]
new_im[:, :, 1] = im[:, :, 1]
new_im[:, :, 2] = im[:, :, 0]
im_list.append(Image.fromarray(new_im.astype('uint8'), 'RGB'))
fc += 1
# Disable gradients
with torch.no_grad():
# Prepare data for model
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
transform_eval = T.Compose([
T.Resize(255, interpolation=Image.BICUBIC),
T.ToTensor(),
T.Normalize(mean, std)
])
model_input = torch.stack([transform_eval(im) for im in im_list])
model_input.to(device)
results, densities = model(model_input)
if len(results) > 1:
results = results.mean()
return results
