''' 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)