''' author ： Hatcher W Email： hcwang@semi.ac.cn use: 密度估计推理 ''' import torch import datetime import random from torch.utils.data import Dataset import numpy as np from network import CSRNet,MCNN,CANNet,PKSNet from dataset import People_Dataset from configs.config import * from configs.config import model_head import json import os from tqdm import tqdm import matplotlib.pyplot as plt import cv2 from torch.utils.data import DataLoader from PIL import Image from PIL import ImageFont,ImageDraw,Image #os.environ['CUDA_VISIBLE_DEVICES'] = '1' device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') result_path = "../data_MPKFNet/result/BirdCounting2" #raw_data_inf_path = '../data_MPKFNet/result_raw/'+"1_3" raw_data_inf_path ="/data1/Hatcher/model/data_MPKFNet/data/BirdCounting/BirdCounting2/train_data/images" display = 1 if not os.path.exists(result_path): os.mkdir(result_path) ''' 准备数据 ''' get_name = [] for _, _, filenames in os.walk(raw_data_inf_path): # root是一个list是正在遍历的文件夹本身的地址 dirs文件夹中的目录 files是文件夹下的所有文件 for filename in filenames: get_name.append(filename) train = raw_data_inf_path trainfile = os.listdir(raw_data_inf_path) train_json = [file for file in trainfile] train_json.sort() with open(inf_json,'w') as f: json.dump(train_json,f) class INf_Dataset(Dataset): def __init__(self, img_list, train=True): random.shuffle(img_list) self.img_list = img_list self.train = train def __len__(self): return len(self.img_list) def __getitem__(self, index): assert index <= len(self), 'index range error' path_temp = raw_data_inf_path + "/" + self.img_list[index] img = cv2.imread(path_temp) # H read image img = cv2.resize(img, (1024, 768)) / 255.0 img = np.transpose(img, [2, 0, 1]) torch.from_numpy(img).type(torch.FloatTensor) return img, self.img_list[index] with open(inf_json, 'r') as outfile: inf_list = json.load(outfile) inf_dataset = INf_Dataset(inf_list) inf_dataset.img_list.sort() data_loader = DataLoader(inf_dataset,1) def valide(): batch_size = 1 with open(trainA_json, 'r') as outfile: train_list = json.load(outfile) with open(testA_json, 'r') as outfile: test_list = json.load(outfile) train_dataset = People_Dataset(train_list) train_loader = DataLoader(train_dataset, batch_size=batch_size) test_dataset = People_Dataset(test_list) test_loader = DataLoader(test_dataset, batch_size=batch_size) net = PKSNet().to(device) checkpoint = torch.load('../data_MPKFNet/weigth/final.pth') net.load_state_dict(checkpoint['net']) print("load model weight ok") # net.load_state_dict(torch.load('../data_MPKFNet/weigth/final.pth')) # print("load model weight ok") mae, mse = 0, 0 if model_head: for img, traget, count, posture, species in tqdm(test_loader): img = img.to(device) target = traget.to(device) t1 = datetime.datetime.now() targets_pred, posture_pred, spice_pred = net(img) t2 = datetime.datetime.now() #print("haoshi ", t2 - t1) mae += abs(targets_pred.data.sum() - target.data.sum()) mse += ((target.data.sum() - targets_pred.data.sum()) * (target.data.sum() - targets_pred.data.sum())) mae = mae / len(test_list) mse = mse / len(test_list) else: for img, traget, countin in tqdm(test_loader): img = img.to(device) target = traget.to(device) t1 = datetime.datetime.now() targets_pred= net(img) t2 = datetime.datetime.now() print("haoshi ", t2 - t1) mae += abs(targets_pred.data.sum() - target.data.sum()) mse += ((target.data.sum() - targets_pred.data.sum()) * (target.data.sum() - targets_pred.data.sum())) mae = mae / len(test_list) mse = mse / len(test_list) print('test_data:{},calculate MAE: \n'.format(len(test_list))) print(' * MAE {mae:.3f} \n'.format(mae=mae)) print(' * MSE {mse:.3f} \n'.format(mse=mse)) def gen_inf(): """ 获取验证数据 :return: """ i, n = 0, len(get_name) while True: if i >= n: i = 0 break batch_x, batch_y ,w_a,h_a=[],[],[],[] path_temp = raw_data_inf_path +"/" +get_name[i] img = cv2.imread(path_temp) # H read image #h, w = np.shape(img)[0], np.shape(img)[1] img = cv2.resize(img, (1024, 768)) / 255.0 #img = img /255.0 img = np.transpose(img, [2, 0, 1]) torch.from_numpy(img).type(torch.FloatTensor) #batch_x.append(img) # batch_y.append(get_name[i]) # w_a.append(w) # h_a.append(h) i += 1 yield img,get_name[i] def inf(devices): net = PKSNet().to(devices) checkpoint = torch.load('../data_MPKFNet/weigth/final.pth') net.load_state_dict(checkpoint['net']) # net.load_state_dict(torch.load('../data_MPKFNet/weigth/sheep/bestmaesheep93,mse29262.pth')) print("load model weight ok") i = 0 for img, img_pathtemp in tqdm(data_loader): img_pathtemp = img_pathtemp[0] img = torch.tensor(img) img = img.to(devices) if model_head: targets_pred, posture_pred, spice_pred = net(img) else: targets_pred = net(img) t1 = datetime.datetime.now() num = targets_pred.data.sum() targets_pred1 = targets_pred.cpu() targets_pred1 = targets_pred1.data.numpy() t2 = datetime.datetime.now() if i % display == 0: num_display = num plt.figure(num_display) #img = cv2.putText(targets_pred1[0, 0, :, :], '3', (5, 5), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (255, 255, 255), 2) # cv2.imshow("hatcher", targets_pred1[0, 0, :, :]) # cv2.waitKey() # cv2.imwrite(result_path + "/"+str(img_pathtemp), targets_pred1[0, 0, :, :]*255) plt.title("predict:{}".format(int(num_display))) plt.imshow(targets_pred1[0, 0, :, :]) plt.savefig(result_path + "/" + str(img_pathtemp)) plt.close() i += 1 # # net = PKSNet().to(device) # checkpoint = torch.load('../data_MPKFNet/weigth/final.pth') # net.load_state_dict(checkpoint['net']) # #net.load_state_dict(torch.load('../data_MPKFNet/weigth/sheep/bestmaesheep93,mse29262.pth')) # print("load model weight ok") # # i=0 # for img, img_pathtemp in tqdm(data_loader): # # img_pathtemp = img_pathtemp[0] # img =torch.tensor(img) # img = img.to(device) # if model_head: # targets_pred, posture_pred, spice_pred = net(img) # else: # targets_pred = net(img) # #cv2.imshow("pre is {}".format(targets_pred.data.sum()),targets_pred.data.numpy()) # # t1 = datetime.datetime.now() # num = targets_pred.data.sum() # targets_pred1 = targets_pred.cpu() # targets_pred1 = targets_pred1.data.numpy() # t2 = datetime.datetime.now() # # print("haoshi ", t2 - t1) # # if i < 343: # # num=0 # if i%display ==0: # num_display = num # fontpath = "font/simsun.ttc" # #font = ImageFont.truetype(fontpath,32) # img = cv2.putText(targets_pred1[0, 0, :, :],'3',(58,5),cv2.FONT_HERSHEY_SIMPLEX,1.2,(255,255,255),2) # #img_pil = Image.fromarray(targets_pred1[0, 0, :, :]) # #draw = ImageDraw.Draw(img_pil) # #draw.text((58, 5), num_display, font=cv2.FONT_HERSHEY_SIMPLEX, fill=(255, 255, 255)) # #bk_img = np.array(img_pil) # cv2.imshow("hatcher", targets_pred1[0, 0, :, :]) # #cv2.waitKey() # # cv2.imwrite(result_path + "/"+str(img_pathtemp), img) # # plt.figure(num_display) # plt.title("predict:{}".format(int(num_display))) # plt.imshow(targets_pred1[0, 0, :, :]) # plt.savefig(result_path + "/"+str(img_pathtemp)) # plt.close() # # Image.fromarray(drawn_image).save(os.path.join(output_dir, image_name)) # i+=1 if __name__ == '__main__': # valide() # valide() inf(device)