238 lines
8.2 KiB
Python
Executable File
238 lines
8.2 KiB
Python
Executable File
'''
|
||
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)
|