#!/usr/bin/python # -*- encoding: utf-8 -*- import torch from matplotlib import pyplot as plt from torch.utils.data import Dataset import torchvision.transforms as transforms import os.path as osp import os from PIL import Image import numpy as np import json import cv2 import time from transform import * class Heliushuju(Dataset): def __init__(self, rootpth, cropsize=(640, 480), mode='train', randomscale=(0.125, 0.25, 0.375, 0.5, 0.675, 0.75, 0.875, 1.0, 1.25, 1.5), *args, **kwargs): super(Heliushuju, self).__init__(*args, **kwargs) assert mode in ('train', 'val', 'test', 'trainval') self.mode = mode print('self.mode', self.mode) self.ignore_lb = 255 with open('./heliushuju_info.json', 'r') as fr: labels_info = json.load(fr) # print('###line30:',labels_info) # self.lb_map = {el['id']: el['trainId'] for el in labels_info} self.lb_map = {el['id']: el['color'] for el in labels_info} # print('###line32:', self.lb_map) # parse img directory self.imgs = {} imgnames = [] impth = osp.join(rootpth, mode, 'images') # 图片所在目录的路径 folders = os.listdir(impth) # 图片名列表 names = [el.replace(el[-4:], '') for el in folders] # el是整个图片名,names是图片名前缀 impths = [osp.join(impth, el) for el in folders] # 图片路径 imgnames.extend(names) # 存放图片名前缀的列表 self.imgs.update(dict(zip(names, impths))) # parse gt directory self.labels = {} gtnames = [] gtpth = osp.join(rootpth, mode, 'labels_2') folders = os.listdir(gtpth) names = [el.replace(el[-4:], '') for el in folders] lbpths = [osp.join(gtpth, el) for el in folders] gtnames.extend(names) self.labels.update(dict(zip(names, lbpths))) self.imnames = imgnames self.len = len(self.imnames) print('self.len', self.mode, self.len) assert set(imgnames) == set(gtnames) assert set(self.imnames) == set(self.imgs.keys()) assert set(self.imnames) == set(self.labels.keys()) # pre-processing self.to_tensor = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), ]) self.trans_train = Compose([ ColorJitter( brightness = 0.5, contrast = 0.5, saturation = 0.5), HorizontalFlip(), # RandomScale((0.25, 0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0)), RandomScale(randomscale), # RandomScale((0.125, 1)), # RandomScale((0.125, 0.25, 0.375, 0.5, 0.675, 0.75, 0.875, 1.0)), # RandomScale((0.125, 0.25, 0.375, 0.5, 0.675, 0.75, 0.875, 1.0, 1.125, 1.25, 1.375, 1.5)), RandomCrop(cropsize) ]) self.mean = (0.485, 0.456, 0.406) self.std = (0.229, 0.224, 0.225) def __getitem__(self, idx): fn = self.imnames[idx] impth = self.imgs[fn] lbpth = self.labels[fn] img = Image.open(impth).convert('RGB') # img = cv2.imread(impth);img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # label = Image.open(lbpth) # 改动 label = cv2.imread(lbpth) # 原始 label = cv2.cvtColor(label, cv2.COLOR_BGR2RGB) # 添加（训练交通事故数据，添加了这行代码使标签颜色正确） # plt.figure(1);plt.imshow(label);plt.show() # 添加 if self.mode == 'train' or self.mode == 'trainval' or self.mode == 'val': label = Image.fromarray(label) im_lb = dict(im = img, lb = label) im_lb = self.trans_train(im_lb) img, label = im_lb['im'], im_lb['lb'] # img = self.to_tensor(img) img = np.array(img); img_bak = img.copy() img = self.preprocess_image(img) label = cv2.resize(np.array(label), (640, 360)) label = label.astype(np.int64)[np.newaxis, :] # 给行上增加维度 # label = cv2.resize(label,(640,360)) # print('###line108:', self.lb_map) label = self.convert_labels(label) # plt.figure(0);plt.imshow(label[0]); # plt.figure(1);plt.imshow(img_bak);plt.show() return img, label.astype(np.int64) def __len__(self): return self.len def convert_labels(self, label): b, h, w, c = label.shape # print('####line118:',label.shape) # b, h, w = label.shape # [1,360,640] label_index = np.zeros((b, h, w)) for k, v in self.lb_map.items(): t_0 = (label[..., 0] == v[0]) t_1 = (label[..., 1] == v[1]) t_2 = (label[..., 2] == v[2]) t_loc = (t_0 & t_1 & t_2) label_index[t_loc] = k # label[label == k] = v # print(label) # print("6666666666666666") return label_index def preprocess_image(self, image): time0 = time.time() image = cv2.resize(image, (640, 360)) time1 = time.time() image = image.astype(np.float32) image /= 255.0 time2 = time.time() # image = image * 3.2 - 1.6 image[:, :, 0] -= self.mean[0] image[:, :, 1] -= self.mean[1] image[:, :, 2] -= self.mean[2] time3 = time.time() image[:, :, 0] /= self.std[0] image[:, :, 1] /= self.std[1] image[:, :, 2] /= self.std[2] time4 = time.time() image = np.transpose(image, (2, 0, 1)) time5 = time.time() image = torch.from_numpy(image).float() # image = image.unsqueeze(0) # outStr = '###line84: in preprocess: resize:%.1f norm:%.1f mean:%.1f std:%.1f trans:%.f ' % ( # self.get_ms(time1, time0), self.get_ms(time2, time1), self.get_ms(time3, time2), self.get_ms(time4, time3), # self.get_ms(time5, time4)) # print(outStr) # print('###line84: in preprocess: resize:%.1f norm:%.1f mean:%.1f std:%.1f trans:%.f '%(self.get_ms(time1,time0),self.get_ms(time2,time1),self.get_ms(time3,time2),self.get_ms(time4,time3) ,self.get_ms(time5,time4) ) ) return image if __name__ == "__main__": from tqdm import tqdm # ds = Heliushuju('./data/', n_classes=2, mode='val') # 原始 ds = Heliushuju('./data/', n_classes=3, mode='val') # 改动 uni = [] for im, lb in tqdm(ds): lb_uni = np.unique(lb).tolist() uni.extend(lb_uni) print(uni) print(set(uni))