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