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