201 lines
7.4 KiB
Python
Executable File
201 lines
7.4 KiB
Python
Executable File
'''
|
||
author ： Hatcher W
|
||
Email： hcwang@semi.ac.cn
|
||
use: 密度估计数据集预处理
|
||
'''
|
||
from torch.utils.data import Dataset
|
||
import random
|
||
import numpy as np
|
||
import cv2
|
||
from scipy.io import loadmat
|
||
import json
|
||
from configs.configs import *
|
||
from configs.configs import model_head
|
||
import matplotlib.pyplot as plt
|
||
import torch
|
||
|
||
class People_Dataset(Dataset):
|
||
def __init__(self, img_list, train=True):
|
||
random.shuffle(img_list)
|
||
self.img_list = img_list
|
||
self.train = train
|
||
|
||
def get_pixels(self, target_path, target_h, target_w):
|
||
mat_data = loadmat(target_path)
|
||
positions = mat_data['image_info'][0, 0][0, 0][0]
|
||
count = mat_data['image_info'][0][0][0][0][1][0][0]
|
||
if model_head:
|
||
posture = mat_data['image_info'][0][0][0][0][2][0][0]
|
||
species = mat_data['image_info'][0][0][0][0][3][0][0]
|
||
if (species >= 0) & (species <= 100):
|
||
species = 0
|
||
if (species >= 101) & (species <= 200):
|
||
species = 1
|
||
if (species >= 201) & (species <= 300):
|
||
species = 2
|
||
if (species >= 301) & (species <= 400):
|
||
species = 3
|
||
if (species >= 401) & (species <= 500):
|
||
species = 4
|
||
if (species >= 501) & (species <= 600):
|
||
species = 5
|
||
|
||
if (species >= 601) & (species <= 700):
|
||
species = 6
|
||
if (species >= 701) & (species <= 800):
|
||
species = 7
|
||
if (species >= 801) & (species <= 900):
|
||
species = 8
|
||
if (species >= 901) & (species <= 1000):
|
||
species = 9
|
||
if (species >= 1001) & (species <= 1100):
|
||
species = 10
|
||
if (species >= 1101) & (species <= 1200):
|
||
species = 11
|
||
|
||
if (species >= 1201):
|
||
species = 12
|
||
pixels = np.zeros((768, 1024))
|
||
for p in positions:
|
||
p[0], p[1] = int(p[0] * 1024 / target_w), int(p[1] * 768 / target_h)
|
||
if p[0] >= 1024 or p[1] >= 768:
|
||
pass
|
||
else:
|
||
pixels[int(p[1]), int(p[0])] += 1
|
||
density = cv2.GaussianBlur(pixels, (15, 15), 0)
|
||
|
||
|
||
if model_head:
|
||
temp = (density, count,posture,species)
|
||
else:
|
||
temp = (density,count)
|
||
return temp
|
||
|
||
def __len__(self):
|
||
return len(self.img_list)
|
||
|
||
def __getitem__(self, index):
|
||
assert index <= len(self), 'index range error'
|
||
img_path = self.img_list[index]
|
||
img_temppath = img_path.replace('../data_MPKFNet/data/test_data/images/','')
|
||
target_path = img_path.replace('.jpg', '.mat').replace('images', 'ground_truth').replace('HatIMG_', 'GT_HatIMG_')
|
||
img = cv2.imread(img_path)
|
||
#print(img_path)
|
||
if img is None:
|
||
print(img_path)
|
||
target_h, target_w = img.shape[0], img.shape[1]
|
||
|
||
img = cv2.resize(img, (1024, 768)) / 255.0
|
||
# img[:, :, 0] = (img[:,:,0] - 0.485) / 0.229
|
||
# img[:, :, 1] = (img[:, :, 0] - 0.456) / 0.224
|
||
# img[:, :, 2] = (img[:, :, 0] - 0.406) / 0.225
|
||
if model_head:
|
||
density, count,posture, species= self.get_pixels(target_path, target_h, target_w)
|
||
else:
|
||
density, count = self.get_pixels(target_path, target_h, target_w)
|
||
#density = cv2.resize(density, (1024//16, 768//16))*256
|
||
density = cv2.resize(density, (1024//8, 768//8))*64
|
||
#density = cv2.resize(density, (1024 // 4, 768 // 4)) *16
|
||
#density = cv2.resize(density, (1024 // 2, 768 // 2)) * 4
|
||
#density = cv2.resize(density, (1024 , 768 ))
|
||
density = np.expand_dims(density, 0)
|
||
img = np.transpose(img, [2, 0, 1])
|
||
if model_head:
|
||
temp = (torch.from_numpy(img).type(torch.FloatTensor), torch.from_numpy(density).type(torch.FloatTensor), torch.tensor(int(count)),\
|
||
torch.tensor(np.array(posture)).type(torch.long),torch.tensor(np.array(species)).type(torch.long))
|
||
else:
|
||
temp = (torch.from_numpy(img).type(torch.FloatTensor), torch.from_numpy(density).type(torch.FloatTensor), torch.tensor(int(count)))
|
||
return temp
|
||
|
||
|
||
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'
|
||
img_path = self.img_list[index]
|
||
img_temppath = img_path.replace('../data_MPKFNet/data/test_data/images/','')
|
||
target_path = img_path.replace('.jpg', '.mat').replace('images', 'ground_truth').replace('HatIMG_', 'GT_HatIMG_')
|
||
img = cv2.imread(img_path)
|
||
|
||
target_h, target_w = img.shape[0], img.shape[1]
|
||
img = cv2.resize(img, (1024, 768)) / 255.0
|
||
# img[:, :, 0] = (img[:,:,0] - 0.485) / 0.229
|
||
# img[:, :, 1] = (img[:, :, 0] - 0.456) / 0.224
|
||
# img[:, :, 2] = (img[:, :, 0] - 0.406) / 0.225
|
||
if model_head:
|
||
density, count, posture, species = self.get_pixels(target_path, target_h, target_w)
|
||
else:
|
||
density, count = self.get_pixels(target_path, target_h, target_w)
|
||
density = cv2.resize(density, (1024//8, 768//8))*64
|
||
#density = cv2.resize(density, (1024,768))
|
||
density = np.expand_dims(density, 0)
|
||
img = np.transpose(img, [2, 0, 1])
|
||
if model_head:
|
||
temp = (torch.from_numpy(img).type(torch.FloatTensor), torch.from_numpy(density).type(torch.FloatTensor),
|
||
torch.tensor(int(count)), \
|
||
torch.tensor(np.array(posture)).type(torch.long), torch.tensor(np.array(species)).type(torch.long))
|
||
else:
|
||
temp = (torch.from_numpy(img).type(torch.FloatTensor), torch.from_numpy(density).type(torch.FloatTensor),
|
||
torch.tensor(int(count)))
|
||
return temp
|
||
|
||
if __name__ == '__main__':
|
||
with open(trainA_json, 'r') as outfile:
|
||
train_list = json.load(outfile)
|
||
count_all =0
|
||
dataset = People_Dataset(train_list)
|
||
if model_head:
|
||
for img, target, count,posture, species in dataset:
|
||
#print(img.shape, target.shape)
|
||
img = img.numpy()
|
||
|
||
target = target.squeeze(0).numpy()
|
||
|
||
count = count.numpy()
|
||
# posture = posture.numpy()
|
||
# species = species.numpy()
|
||
# print(posture)
|
||
# print(species)
|
||
# print("posture is {} ,species is {}".format(posture,species))
|
||
img = np.transpose(img, [1,2,0])
|
||
#
|
||
# plt.figure('img')
|
||
# plt.imshow((img[:, :, ::-1]*255).astype(np.uint8))
|
||
# plt.show()
|
||
# plt.figure('target')
|
||
# plt.imshow(target)
|
||
# plt.show()
|
||
print(abs(target.sum() - count))
|
||
count_all =count_all+ abs(target.sum() - count)
|
||
print(count_all)
|
||
# print(target.sum() - count)
|
||
else:
|
||
for img, target, count in dataset:
|
||
print(img.shape, target.shape)
|
||
img = img.numpy()
|
||
|
||
target = target.squeeze(0).numpy()
|
||
|
||
count = count.numpy()
|
||
|
||
img = np.transpose(img, [1, 2, 0])
|
||
#
|
||
plt.figure('img')
|
||
plt.imshow((img[:, :, ::-1] * 255).astype(np.uint8))
|
||
plt.show()
|
||
plt.figure('target')
|
||
plt.imshow(target)
|
||
plt.show()
|
||
print(target.sum() - count)
|
||
|
||
|
||
|
||
|