def main():
global args
global support_imgs, support_gt, support_imgs_fea
args = parser.parse_args()
with open(args.test_json, 'r') as outfile:
val_list = json.load(outfile)
residual = Residual(k_nn=args.n)
residual = residual.cuda()
counter = CSRNet()
counter = counter.cuda()
# load counter params
checkpoint = torch.load('./saved_models/countercheckpoint.pth.tar')
counter.load_state_dict(checkpoint['state_dict_model'])
# load residual regressor params
checkpoint = torch.load('./saved_models/residualcheckpoint.pth.tar')
residual.load_state_dict(checkpoint['state_dict_res'])
support_imgs = checkpoint['support_imgs'].cuda()
support_gt = checkpoint['support_gt'].cuda()
counter(support_imgs)
support_imgs_fea = counter.features
app, res, final = validate(val_list, counter, residual)
def main():
transform = ST.Compose(
[
ST.ToNumpyForVal(),
ST.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
]
)
global args
args = parser.parse_args()
model = CSRNet()
model = model.to("cuda")
# checkpoint = flow.load('checkpoint/Shanghai_BestModelA/shanghaiA_bestmodel')
checkpoint = flow.load(args.modelPath)
model.load_state_dict(checkpoint)
img = transform(Image.open(args.picPath).convert("RGB"))
img = flow.Tensor(img)
img = img.to("cuda")
output = model(img.unsqueeze(0))
print("Predicted Count : ", int(output.detach().to("cpu").sum().numpy()))
temp = output.view(output.shape[2], output.shape[3])
temp = temp.numpy()
plt.title("Predicted Count")
plt.imshow(temp, cmap=c.jet)
plt.show()
temp = h5py.File(args.picDensity, "r")
temp_1 = np.asarray(temp["density"])
plt.title("Original Count")
plt.imshow(temp_1, cmap=c.jet)
print("Original Count : ", int(np.sum(temp_1)) + 1)
plt.show()
print("Original Image")
plt.title("Original Image")
plt.imshow(plt.imread(args.picPath))
plt.show()
def __init__(self):
self.best_pred = 1e6
# Define Saver
self.saver = Saver(opt)
self.saver.save_experiment_config()
# visualize
if opt.visualize:
# vis_legend = ["Loss", "MAE"]
# batch_plot = create_vis_plot(vis, 'Batch', 'Loss', 'batch loss', vis_legend[0:1])
# val_plot = create_vis_plot(vis, 'Epoch', 'result', 'val result', vis_legend[1:2])
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
# Dataset dataloader
self.train_dataset = SHTDataset(opt.train_dir, train=True)
self.train_loader = DataLoader(self.train_dataset,
num_workers=opt.workers,
shuffle=True,
batch_size=opt.batch_size) # must be 1
self.test_dataset = SHTDataset(opt.test_dir, train=False)
self.test_loader = torch.utils.data.DataLoader(
self.test_dataset, shuffle=False, batch_size=opt.batch_size
) # must be 1, because per image size is different
torch.cuda.manual_seed(opt.seed)
model = CSRNet()
self.model = model.to(opt.device)
if opt.resume:
if os.path.isfile(opt.pre):
print("=> loading checkpoint '{}'".format(opt.pre))
checkpoint = torch.load(opt.pre)
opt.start_epoch = checkpoint['epoch']
self.best_pred = checkpoint['best_pred']
self.model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})".format(
opt.pre, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(opt.pre))
if opt.use_mulgpu:
self.model = torch.nn.DataParallel(self.model,
device_ids=opt.gpu_id)
self.criterion = nn.MSELoss(reduction='mean').to(opt.device)
self.optimizer = torch.optim.SGD(self.model.parameters(),
opt.lr,
momentum=opt.momentum,
weight_decay=opt.decay)
# Define lr scheduler
self.scheduler = lr_scheduler.MultiStepLR(
self.optimizer,
milestones=[round(opt.epochs * x) for x in opt.steps],
gamma=opt.scales)
self.scheduler.last_epoch = opt.start_epoch - 1
def __init__(self, checkpoint_path='0model_best.pth.tar'):
self.transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
self.model = CSRNet()
self.model = self.model.cuda()
self.checkpoint = torch.load(checkpoint_path)
self.model.load_state_dict(self.checkpoint['state_dict'])
def load_model(model_path, device):
"""
model_path: saved model (.pth or .pth.tar)
#TODO: map_location
"""
model = CSRNet()
checkpoint = torch.load(model_path, map_location=device)
model.load_state_dict(checkpoint['state_dict'])
return model
def initialize_model():
transform = transforms.Compose(
[
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
]
)
model = CSRNet().cuda()
checkpoint = torch.load(path.join("checkpoints", "model_best.pth.tar"))
model.load_state_dict(checkpoint["state_dict"])
return transform, model
def main():
global args
print(args.path_testing_image)
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
img_paths = []
for img_path in glob.glob(os.path.join(args.path_testing_image, '*.png')):
img_paths.append(img_path)
model = CSRNet()
#defining the model
model = model.cuda()
#loading the trained weights
checkpoint = torch.load(args.best_model_csrnet_path)
model.load_state_dict(checkpoint['state_dict'])
img = transform(Image.open(args.path_testing_image).convert('RGB')).cuda()
output = model(img.unsqueeze(0))
print("Predicted Count : ", int(output.detach().cpu().sum().numpy()))
temp = np.asarray(output.detach().cpu().reshape(
output.detach().cpu().shape[2],
output.detach().cpu().shape[3]))
plt.imshow(temp, cmap=c.jet)
plt.axis('off')
plt.savefig("predicted_dt.png", bbox_inches='tight')
temp = h5py.File(args.path_testing_image.replace('.png', '.h5'))
temp_1 = np.asarray(temp['density'])
plt.imshow(temp_1, cmap=c.jet)
print(" Original Count : ", int(np.sum(temp_1)) + 1)
plt.axis('off')
plt.savefig("original_dt.png", bbox_inches='tight')
print("Original Image")
plt.imshow(plt.imread(args.path_testing_image))
plt.axis('off')
plt.savefig("original_image.png", bbox_inches='tight')
f = open('results.txt', 'w')
f.write("Predicted Count : " +
str(int(output.detach().cpu().sum().numpy())) +
" Original Count : " + str(int(np.sum(temp_1)) + 1))
f.close()
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 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 prediction(path):
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
img = transform(Image.open(path).convert('RGB')).cuda()
model_best = CSRNet()
model_best = model_best.cuda()
checkpoint = torch.load('model_best.pth.tar')
model_best.load_state_dict(checkpoint['state_dict'])
output = model_best(img.unsqueeze(0))
pred = int(output.detach().cpu().sum().numpy())
fake_pred1 = random.randint(0, 100) + pred
fake_pred2 = pred - random.randint(0, 50)
return pred, fake_pred1, fake_pred2
class PeopleCounter:
def __init__(self, checkpoint_path='0model_best.pth.tar'):
self.transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
self.model = CSRNet()
self.model = self.model.cuda()
self.checkpoint = torch.load(checkpoint_path)
self.model.load_state_dict(self.checkpoint['state_dict'])
def countPeople(self, img):
img = self.transform(img.convert('RGB')).cuda()
output = self.model(img.unsqueeze(0))
return int(output.detach().cpu().sum().numpy())
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 main():
transform = ST.Compose([
ST.ToNumpyForVal(),
ST.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
global args
args = parser.parse_args()
root = "./dataset/"
# now generate the ShanghaiA's ground truth
part_A_train = os.path.join(root, "part_A_final/train_data", "images")
part_A_test = os.path.join(root, "part_A_final/test_data", "images")
part_B_train = os.path.join(root, "part_B_final/train_data", "images")
part_B_test = os.path.join(root, "part_B_final/test_data", "images")
path_sets = []
if args.picSrc == "part_A_test":
path_sets = [part_A_test]
elif args.picSrc == "part_B_test":
path_sets = [part_B_test]
img_paths = []
for path in path_sets:
for img_path in glob.glob(os.path.join(path, "*.jpg")):
img_paths.append(img_path)
model = CSRNet()
model = model.to("cuda")
checkpoint = flow.load(args.modelPath)
model.load_state_dict(checkpoint)
MAE = []
for i in range(len(img_paths)):
img = transform(Image.open(img_paths[i]).convert("RGB"))
img = np.asarray(img).astype(np.float32)
img = flow.Tensor(img, dtype=flow.float32, device="cuda")
img = img.to("cuda")
gt_file = h5py.File(
img_paths[i].replace(".jpg",
".h5").replace("images", "ground_truth"), "r")
groundtruth = np.asarray(gt_file["density"])
with flow.no_grad():
output = model(img.unsqueeze(0))
mae = abs(output.sum().numpy() - np.sum(groundtruth))
MAE.append(mae)
avg_MAE = sum(MAE) / len(MAE)
print("test result: MAE:{:2f}".format(avg_MAE))
def main():
global args, best_prec1
best_prec1 = 1e6
args = parser.parse_args() # 将变量以标签-真值的字典形式存入args字典中
args.original_lr = 1e-7
args.lr = 1e-7
args.batch_size = 1
args.momentum = 0.95
args.decay = 5 * 1e-4
args.start_epoch = 0
args.epochs = 400
args.steps = [-1, 1, 100, 150]
args.scales = [1, 1, 1, 1]
args.workers = 4
args.seed = time.time()
args.print_freq = 30
with open(args.train_json, 'r') as outfile: # 打开训练集.json文件
train_list = json.load(outfile) ##得到训练集图片全部的地址，里面是以列表形式存储的所有图片的绝对路径
with open(args.test_json, 'r') as outfile: # 打开测试集.json文件
val_list = json.load(outfile) #得到验证集图片的全部地址
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
torch.cuda.manual_seed(args.seed)
model = CSRNet()
model = model.cuda()
criterion = nn.MSELoss(size_average=False).cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.decay)
if args.pre: ##预训练模型给出
if os.path.isfile(args.pre):
print("=> loading checkpoint '{}'".format(args.pre))
checkpoint = torch.load(args.pre)
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.pre, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.pre))
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch)
train(train_list, model, criterion, optimizer, epoch)
prec1 = validate(val_list, model, criterion)
is_best = prec1 < best_prec1
best_prec1 = min(prec1, best_prec1) #mae的比较
print(' * best MAE {mae:.3f} '.format(mae=best_prec1))
save_checkpoint({ ##保存
'epoch': epoch + 1,
'arch': args.pre,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer' : optimizer.state_dict(),
}, is_best,args.task)
def count(img):
model = CSRNet()
if C["cuda"]:
model = model.cuda()
pth = torch.load(C["pth"])
else:
model = model.cpu()
pth = torch.load(C["pth"], map_location="cpu")
model.load_state_dict(pth["state_dict"])
# img = 255.0 * to_tensor(Image.open(img_path).convert("RGB"))
# for i in range(3):
# img[i,:,:] = img[i,:,:] + C["img_corr"][i]
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=C["mean"], std=C["std"])
])
img = transform(img.convert("RGB"))
if C["cuda"]:
img = img.cuda()
else:
img = img.cpu()
output = model(img.unsqueeze(0)).detach().cpu()
dmap = np.asarray(output.reshape(output.shape[2], output.shape[3]))
count = int(np.sum(dmap))
return count
def count2(img_root, model_param_path):
# device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
device = torch.device("cpu")
model = CSRNet()
model.load_state_dict(torch.load(model_param_path))
model.to(device)
test_dataloader = create_test_dataloader(img_root) # dataloader
for i, data in enumerate(tqdm(test_dataloader, ncols=50)):
image = data['image'].to(device)
et_dmp = model(image).detach()
# count = et_densitymap.data.sum()
#count = str('%.2f' % (et_densitymap[0].cpu().sum()))
# et_dmp = et_densitymap[0] / torch.max(et_densitymap[0])
et_dmp = et_dmp.numpy()
et_dmp = et_dmp[0][0]
count = np.sum(et_dmp)
plt.figure(i)
plt.axis("off")
plt.imshow(et_dmp, cmap=CM.jet)
# 去除坐标轴
plt.gca().xaxis.set_major_locator(plt.NullLocator())
plt.gca().yaxis.set_major_locator(plt.NullLocator())
# 输出图片边框设置
plt.subplots_adjust(top=1,
bottom=0,
left=0,
right=1,
hspace=0,
wspace=0)
plt.margins(0, 0)
plt.savefig(img_root + "/test_data/result/" + str(i + 1) + "_dmp" +
".jpg")
print(str(i + 1) + "_" + "renshu:", count)
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 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 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 one_count(img_path, model_param_path):
filename = img_path.split('/')[-1]
filenum = filename.split('.')[0]
save_path = './data_test/test_data/result/' + filenum
if not os.path.exists(save_path):
os.makedirs(save_path)
img_src_save_path = save_path + '/' + filenum + '_src' + '.jpg'
img_et_save_path = img_src_save_path.replace('src', 'et')
img_overlap_save_path = img_src_save_path.replace('src', 'overlap')
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model = CSRNet()
model.load_state_dict(torch.load(model_param_path))
model.to(device)
img_src = open_img(img_path)
img = open_img(img_path)
img_trans = Compose([ToTensor(), Normalize(mean=[0.5, 0.5, 0.5], std=[0.225, 0.225, 0.225])])
img = img_trans(img)
img = Variable(torch.unsqueeze(img, dim=0).float(), requires_grad=False)
# print('img_src size:', img.shape)
et_dmap = model(img)
et_dmap = et_dmap.detach().numpy()
et_dmap = et_dmap[0][0]
people_num = np.sum(et_dmap)
# print(et_dmap.shape)
print(filenum + '_num:', '\t', people_num)
# img_src
plt.figure(0)
plt.imshow(img_src)
plt.axis('off')
# 去除坐标轴
plt.gca().xaxis.set_major_locator(plt.NullLocator())
plt.gca().yaxis.set_major_locator(plt.NullLocator())
# 输出图片边框设置
plt.subplots_adjust(top=1, bottom=0, left=0, right=1, hspace=0, wspace=0)
plt.margins(0, 0)
plt.savefig(img_src_save_path, bbox_inches='tight', dpi=100, pad_inches=-0.04)
#
# # img_et
plt.figure(1)
plt.imshow(et_dmap, cmap=CM.jet)
plt.axis('off')
# 去除坐标轴
plt.gca().xaxis.set_major_locator(plt.NullLocator())
plt.gca().yaxis.set_major_locator(plt.NullLocator())
# 输出图片边框设置
plt.subplots_adjust(top=1, bottom=0, left=0, right=1, hspace=0, wspace=0)
plt.margins(0, 0)
plt.savefig(img_et_save_path, bbox_inches='tight', dpi=100, pad_inches=-0.04)
#
img_src = cv2.imread(img_src_save_path)
img_et = cv2.imread(img_et_save_path)
img_et = cv2.resize(img_et, (img_src.shape[1], img_src.shape[0]))
img_overlap = cv2.addWeighted(img_src, 0.2, img_et, 0.8, 0)
cv2.imwrite(img_overlap_save_path, img_overlap)
return people_num
def count1(img_root, model_param_path):
device = torch.device(
"cuda") if torch.cuda.is_available() else torch.device("cpu")
model = CSRNet()
model.load_state_dict(torch.load(model_param_path))
model.to(device)
test_dataloader = create_test_dataloader(img_root) # dataloader
# 添加进度条
for i, data in enumerate(tqdm(test_dataloader, ncols=50)):
image = data['image'].to(device)
et_densitymap = model(image).detach()
# count = et_densitymap.data.sum()
count = str('%.2f' % (et_densitymap[0].cpu().sum()))
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)))
def count3(img_root, model_param_path):
writer = SummaryWriter()
device = torch.device(
"cuda") if torch.cuda.is_available() else torch.device("cpu")
model = CSRNet()
model.load_state_dict(torch.load(model_param_path))
model.to(device)
test_dataloader = create_test_dataloader(img_root) # dataloader
for i, data in enumerate(tqdm(test_dataloader, ncols=50)):
image = data['image'].to(device)
et_densitymap = model(image).detach()
# count = et_densitymap.data.sum()
count = str('%.2f' % (et_densitymap[0].cpu().sum()))
writer.add_image(str(i) + '/img：', denormalize(image[0].cpu()))
writer.add_image(
str(i) + "/dmp_count:" + count,
et_densitymap[0] / torch.max(et_densitymap[0]))
print(str(i + 1) + "_img count success")
def run_test(model_path_1, model_path_2, test):
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
root = os.path.dirname(os.path.abspath(__file__))
# In[4]:
very_small = '<10' if test == 'B' else '<200'
small = '10-50' if test == 'B' else '200-500'
medium = '50-100' if test == 'B' else '500-1000'
large = '100-200' if test == 'B' else '1000-1500'
very_large = '200<' if test == 'B' else '1500<'
very_small_criterion = 10 if test == 'B' else 200
small_criterion = 50 if test == 'B' else 500
medium_criterion = 100 if test == 'B' else 1000
large_criterion = 200 if test == 'B' else 1500
result_dict_mea = {
very_small: [],
small: [],
medium: [],
large: [],
very_large: []
}
result_dict_msa = {
very_small: [],
small: [],
medium: [],
large: [],
very_large: []
}
result_dict_msa_pixel = {
very_small: [],
small: [],
medium: [],
large: [],
very_large: []
}
result_dict_psnr = {
very_small: [],
small: [],
medium: [],
large: [],
very_large: []
}
result_dict_gt = {
very_small: [],
small: [],
medium: [],
large: [],
very_large: []
}
result_dict_swd = {
very_small: [],
small: [],
medium: [],
large: [],
very_large: []
}
result_dict_mea_1 = {
very_small: [],
small: [],
medium: [],
large: [],
very_large: []
}
result_dict_msa_1 = {
very_small: [],
small: [],
medium: [],
large: [],
very_large: []
}
result_dict_msa_pixel_1 = {
very_small: [],
small: [],
medium: [],
large: [],
very_large: []
}
result_dict_psnr_1 = {
very_small: [],
small: [],
medium: [],
large: [],
very_large: []
}
result_dict_gt_1 = {
very_small: [],
small: [],
medium: [],
large: [],
very_large: []
}
result_dict_swd_1 = {
very_small: [],
small: [],
medium: [],
large: [],
very_large: []
}
# now generate the ShanghaiA's ground truth
part_A_train = os.path.join(root, 'part_A_final/train_data', 'images')
part_A_test = os.path.join(root, 'part_A_final/test_data', 'images')
part_B_train = os.path.join(root, 'part_B_final/train_data', 'images')
part_B_test = os.path.join(root, 'part_B_final/test_data', 'images')
path_sets = [part_B_test] if test == 'B' else [part_A_test]
model_name_1 = model_path_1
model_name_2 = model_path_2
img_paths = []
for path in path_sets:
for img_path in glob.glob(os.path.join(path, '*.jpg')):
img_paths.append(img_path)
model_1 = CSRNet()
model_1 = model_1.cuda()
checkpoint = torch.load(model_name_1)
model_1.load_state_dict(checkpoint['state_dict'])
model_2 = CSRNet()
model_2 = model_2.cuda()
checkpoint = torch.load(model_name_2)
model_2.load_state_dict(checkpoint['state_dict'])
from matplotlib import cm
image_list = []
gt_list = []
predict_list_1 = []
predict_list_2 = []
mse_list = []
mae_list = []
psnr_list = []
gt_number_list = []
predict_number_list = []
swd_list = []
original_gt_list = []
correct_resize_gt_list = []
reg_resize_list = []
original_image_list = []
for i in range(len(img_paths)):
plane_image = Image.open(img_paths[i]).convert('RGB')
img = transform(plane_image).cuda()
gt_file = h5py.File(
img_paths[i].replace('.jpg', '.h5').replace('images', 'ground'),
'r')
groundtruth = np.asarray(gt_file['density'])
sum_convovled_kernel = np.ones((8, 8))
target = sg.convolve2d(groundtruth,
sum_convovled_kernel[::-1, ::-1],
mode='valid')[::8, ::8]
output_turch_1 = model_1(img.unsqueeze(0))
output_1 = np.array(output_turch_1.data.cpu()[0, 0, :, :])
cur_mae_1 = abs(output_1.sum() - np.sum(groundtruth))
cur_mse_1 = np.square(output_1.sum() - target.sum())
cur_image_mse_1 = np.square(output_1 - target).sum() / output_1.size
gt_sum = np.sum(groundtruth)
output_turch_2 = model_2(img.unsqueeze(0))
output_2 = np.array(output_turch_2.data.cpu()[0, 0, :, :])
cur_mae_2 = abs(output_2.sum() - np.sum(groundtruth))
cur_mse_2 = np.square(output_2.sum() - target.sum())
cur_image_mse_2 = np.square(output_2 - target).sum() / (
output_2.shape[0] * output_2.shape[1])
cur_psnr_1 = compare_psnr(target, output_1, data_range=1.0)
cur_ssim_1 = compare_ssim(target, output_1)
target_turch = torch.from_numpy(target)
target_turch = target_turch.type(
torch.FloatTensor).unsqueeze(0).unsqueeze(0).cuda()
target_turch = Variable(target_turch)
cur_swd_1 = swd.swd(target_turch,
output_turch_1).detach().float().unsqueeze(
0).data[0].float().item()
resize_output_1 = cv2.resize(
output_1, (int(groundtruth.shape[1]), int(groundtruth.shape[0])),
interpolation=cv2.INTER_CUBIC) / 64
resize_output_2 = cv2.resize(
output_2, (int(groundtruth.shape[1]), int(groundtruth.shape[0])),
interpolation=cv2.INTER_CUBIC) / 64
resize_target = cv2.resize(
target, (int(groundtruth.shape[1]), int(groundtruth.shape[0])),
interpolation=cv2.INTER_CUBIC) / 64
cur_psnr_2 = compare_psnr(target, output_2, data_range=1.0)
cur_swd_2 = swd.swd(target_turch,
output_turch_2).detach().float().unsqueeze(
0).data[0].float().item()
original_image_list.append(np.array(plane_image).astype(np.uint8))
original_gt_list.append(groundtruth)
correct_resize_gt_list.append(target)
reg_resize_list.append(resize_target)
image_list.append(np.array(plane_image).astype(np.uint8))
gt_list.append(resize_target)
predict_list_1.append(resize_output_1)
predict_list_2.append(resize_output_2)
mae_list.append((round(cur_mae_1, 2), round(cur_mae_2, 2)))
mse_list.append((round(cur_mse_1, 2), round(cur_mse_2, 2)))
psnr_list.append((round(cur_psnr_1, 2), round(cur_psnr_2, 2)))
swd_list.append((round(cur_swd_1, 2), round(cur_swd_2, 2)))
gt_number_list.append(round(np.sum(groundtruth), 2))
predict_number_list.append((round(output_1.sum(),
2), round(output_2.sum(), 2)))
run_it = True
if gt_sum <= very_small_criterion:
criterion = very_small
elif gt_sum > very_small_criterion and gt_sum <= small_criterion:
criterion = small
elif gt_sum > small_criterion and gt_sum <= medium_criterion:
criterion = medium
elif gt_sum > medium_criterion and gt_sum <= large_criterion:
criterion = large
else:
criterion = very_large
result_dict_swd[criterion].append(cur_swd_1)
result_dict_gt[criterion].append(gt_sum)
result_dict_mea[criterion].append(cur_mae_1)
result_dict_msa[criterion].append(cur_mse_1)
result_dict_msa_pixel[criterion].append(cur_image_mse_1)
result_dict_psnr[criterion].append(cur_psnr_1)
result_dict_swd_1[criterion].append(cur_swd_2)
result_dict_gt_1[criterion].append(gt_sum)
result_dict_mea_1[criterion].append(cur_mae_2)
result_dict_msa_1[criterion].append(cur_mse_2)
result_dict_msa_pixel_1[criterion].append(cur_image_mse_2)
result_dict_psnr_1[criterion].append(cur_psnr_2)
if run_it and len(image_list) > 4:
#image_list, gt_list, predict_list_1, predict_list_2, mae_list, psnr_list, gt_number_list, predict_number_list,swd_list
fig = create_heat_map_compare_two_models(
image_list, gt_list, predict_list_1, predict_list_2, mae_list,
mse_list, psnr_list, gt_number_list, predict_number_list,
swd_list)
name_1 = model_name_1.split('/')[-1].split('.')[0]
name_2 = model_name_2.split('/')[-1].split('.')[0]
path = os.path.join(root, 'compare_models', name_1 + '_' + name_2)
os.makedirs(path, exist_ok=True)
fig.savefig(
os.path.join(
path, name_1 + '_' + name_2 + '_compare_models_' + str(i) +
'.png'))
# fig_mea = crate_graph_per_number_of_gt(deepcopy(result_dict_mea),'MEA')
# fig_msa = crate_graph_per_number_of_gt(deepcopy(result_dict_msa), 'MSA')
# fig_ssim = crate_graph_per_number_of_gt(deepcopy(result_dict_ssim), 'SSIM')
# fig_psnr = crate_graph_per_number_of_gt(deepcopy(result_dict_psnr), 'PSNR')
# fig_swd = crate_graph_per_number_of_gt(deepcopy(result_dict_swd), 'SWD')
image_list = []
gt_list = []
predict_list_1 = []
predict_list_2 = []
ssim_list = []
mae_list = []
psnr_list = []
gt_number_list = []
predict_number_list = []
print('the psnr_1')
print(i, cur_psnr_1)
print('the psnr_2')
print(i, cur_psnr_2)
print('the mae_1')
print(i, cur_mae_1)
print('the mae_2')
print(i, cur_mae_2)
print('the mse_1')
print(i, cur_mse_1)
print('the mse_2')
print(i, cur_mse_2)
print('the mse_pixel_1')
print(i, cur_image_mse_1)
print('the mse_pixel_1')
print(i, cur_image_mse_2)
name_1 = model_name_1.split('/')[-1].split('.')[0]
name_2 = model_name_2.split('/')[-1].split('.')[0]
cur_path = os.path.join(root, 'compare_models', name_1 + '_' + name_2)
os.makedirs(cur_path, exist_ok=True)
cur_result_dict_mea_1 = {}
cur_result_dict_msa_1 = {}
cur_result_dict_mea_2 = {}
cur_result_dict_msa_2 = {}
for key in result_dict_mea.keys():
cur_result_dict_mea_1[key] = np.sum(result_dict_mea[key]) / np.sum(
result_dict_gt[key]) if np.sum(result_dict_gt[key]) > 0 else 0
cur_result_dict_msa_1[key] = result_dict_msa[key]
cur_result_dict_mea_2[key] = np.sum(result_dict_mea_1[key]) / np.sum(
result_dict_gt[key]) if np.sum(result_dict_gt[key]) > 0 else 0
cur_result_dict_msa_2[key] = result_dict_msa_1[key]
fig_mea = create_graph_per_number_of_gt_vs_two_models(
deepcopy(cur_result_dict_mea_1), deepcopy(cur_result_dict_mea_2),
'MAE')
fig_msa = create_graph_per_number_of_gt_vs_two_models(
deepcopy(cur_result_dict_msa_1), deepcopy(cur_result_dict_msa_2),
'MSE')
# fig_ssim = create_graph_per_number_of_gt_vs_two_models(deepcopy(result_dict_ssim),deepcopy(cur_result_dict_mea), 'SSIM')
fig_psnr = create_graph_per_number_of_gt_vs_two_models(
deepcopy(result_dict_psnr), deepcopy(result_dict_psnr_1), 'PSNR')
fig_swd = create_graph_per_number_of_gt_vs_two_models(
deepcopy(result_dict_swd), deepcopy(result_dict_swd_1), 'SWD')
fig_mse_pixel = create_graph_per_number_of_gt_vs_two_models(
deepcopy(result_dict_msa_pixel), deepcopy(result_dict_msa_pixel_1),
'Pixel MSE')
os.makedirs(cur_path, exist_ok=True)
fig_mea.savefig(os.path.join(cur_path, 'MEA_bar_' + str(i) + '.png'))
os.makedirs(cur_path, exist_ok=True)
fig_msa.savefig(os.path.join(cur_path, 'MSA_bar_' + str(i) + '.png'))
os.makedirs(cur_path, exist_ok=True)
fig_psnr.savefig(os.path.join(cur_path, 'PSNR_bar_' + str(i) + '.png'))
os.makedirs(cur_path, exist_ok=True)
fig_swd.savefig(os.path.join(cur_path, 'swd_bar_' + str(i) + '.png'))
os.makedirs(cur_path, exist_ok=True)
fig_mse_pixel.savefig(
os.path.join(cur_path, 'per_pixel_mse_bar_' + str(i) + '.png'))
def main():
global args, best_prec1
args = make_meow_args()
best_prec1 = 1e6
args.original_lr = 1e-7
args.lr = 1e-7
args.batch_size = 1
args.momentum = 0.95
args.decay = 5 * 1e-4
args.start_epoch = 0
args.epochs = 400
args.steps = [-1, 1, 100, 150]
args.scales = [1, 1, 1, 1]
args.workers = 4
args.seed = time.time()
args.print_freq = 30
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
torch.cuda.manual_seed(args.seed)
DATA_PATH = "/data/cv_data/shanghaitech-with-people-density-map/ShanghaiTech/part_A/train_data"
train_list, val_list = get_train_val_list(DATA_PATH)
model = CSRNet()
model = model.cuda()
criterion = nn.MSELoss(size_average=False).cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.decay)
if args.pre:
if os.path.isfile(args.pre):
print("=> loading checkpoint '{}'".format(args.pre))
checkpoint = torch.load(args.pre)
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.pre, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.pre))
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch)
train(train_list, model, criterion, optimizer, epoch)
prec1 = validate(val_list, model, criterion)
is_best = prec1 < best_prec1
best_prec1 = min(prec1, best_prec1)
print(' * best MAE {mae:.3f} '.format(mae=best_prec1))
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.pre,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
}, is_best, args.task)
import os
import h5py
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(
def main():
global args, best_prec1
best_prec1 = 1e6
args = parser.parse_args()
args.original_lr = 1e-7
args.lr = 1e-7
args.batch_size = 1
args.momentum = 0.95
args.decay = 5 * 1e-4
args.start_epoch = 0
args.epochs = 400
args.steps = [-1, 1, 100, 150]
args.scales = [1, 1, 1, 1]
args.workers = 4
args.seed = time.time()
args.print_freq = 30
# with open(args.train_json, 'r') as outfile:
# train_list = json.load(outfile)
train_list = [
os.path.join(args.train_path, i) for i in os.listdir(args.train_path)
]
# with open(args.test_json, 'r') as outfile:
# val_list = json.load(outfile)
print(train_list)
val_list = [
os.path.join(args.train_path, j) for j in os.listdir(args.test_path)
]
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
torch.cuda.manual_seed(args.seed)
model = CSRNet()
model = model.to(device)
criterion = nn.MSELoss(size_average=False).to(device)
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.decay)
if args.pre:
if os.path.isfile(args.pre):
print("=> loading checkpoint '{}'".format(args.pre))
checkpoint = torch.load(args.pre)
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.pre, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.pre))
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch)
train(train_list, model, criterion, optimizer, epoch)
prec1 = validate(val_list, model, criterion)
is_best = prec1 < best_prec1
best_prec1 = min(prec1, best_prec1)
print(' * best MAE {mae:.3f} '.format(mae=best_prec1))
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.pre,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
}, is_best, args.task)
path_sets = [part_A_test]
# In[5]:
img_paths = []
for path in path_sets:
for img_path in glob.glob(os.path.join(path, '*.jpg')):
img_paths.append(img_path)
# In[6]:
model = CSRNet()
# In[7]:
model = model.cuda()
# In[38]:
checkpoint = torch.load('0model_best.pth.tar')
# In[39]:
--train_image_root /home/gohyojun/바탕화면/Anthroprocene/Dataset/crane
--train_image_gt_root /home/gohyojun/바탕화면/Anthroprocene/Dataset/crane_labeled
--train_image_density_root /home/gohyojun/바탕화면/Anthroprocene/Dataset/density_map
--test_image_root /home/gohyojun/바탕화면/Anthroprocene/Dataset/crane
--test_image_gt_root /home/gohyojun/바탕화면/Anthroprocene/Dataset/crane_labeled
--test_image_density_root /home/gohyojun/바탕화면/Anthroprocene/Dataset/density_map
"""
if __name__=="__main__":
# argument parsing.
args = parser.parse_args()
cfg = Config(args) # configuration
model = CSRNet().to(cfg.device) # model
criterion = nn.MSELoss(size_average=False) # objective
optimizer = torch.optim.Adam(model.parameters(),lr=cfg.lr) # optimizer
train_dataloader = create_train_dataloader(cfg.train_dataset_root, use_flip=True, batch_size=cfg.batch_size)
test_dataloader = create_test_dataloader(cfg.test_dataset_root) # dataloader
min_mae = sys.maxsize
min_mae_epoch = -1
for epoch in range(1, cfg.epochs): # start training
model.train()
epoch_loss = 0.0
for i, data in enumerate(tqdm(train_dataloader)):
image = data['image'].to(cfg.device)
gt_densitymap = data['densitymap'].to(cfg.device) * 16# todo 1/4 rescale effect때문에
et_densitymap = model(image) # forward propagation
def denormalize(tensor):
mean = [0.5, 0.5, 0.5]
std = [0.225, 0.225, 0.225]
for t, m, s in zip(tensor, mean, std):
t.mul_(s).add_(m)
return tensor
if __name__ == "__main__":
cfg = Config() # configuration
continue_training = False
model = CSRNet().to(cfg.device)
if continue_training:
model.load_state_dict(
torch.load('checkpoints/shaghai_tech_a_best.pth')) # GPU
#torch.load('checkpoints/shaghai_tech_a_best.pth', map_location=lambda storage, loc: storage) # CPU
# model
criterion = nn.MSELoss(size_average=False) # objective
optimizer = torch.optim.Adam(model.parameters(), lr=cfg.lr) # optimizer
train_dataloader = create_train_dataloader(cfg.dataset_root,
use_flip=True,
batch_size=cfg.batch_size)
test_dataloader = create_test_dataloader(cfg.dataset_root) # dataloader
min_mae = sys.maxsize
min_mae_epoch = -1
