first commit
This commit is contained in:
Semi212 committed
commit
54ccfb11fe
218 files changed
+19525
No files matched your search
Executable
+237
@@ -0,0 +1,237 @@
|
||||
'''
|
||||
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)
|
||||
Reference in new issue
Block a user