198 lines
6.5 KiB
Python
198 lines
6.5 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
@project:
|
||
@File : def_selfv5
|
||
@Author : qiqq
|
||
@create_time : 2023/3/22 22:16
|
||
在my的基础上aug后的
|
||
"""
|
||
from __future__ import print_function, division
|
||
import os
|
||
from PIL import Image
|
||
import matplotlib.pyplot as plt
|
||
from torch.utils.data import Dataset
|
||
from taihuyuan_roof.dataloaders.mypath import Path #注意一下这个path
|
||
from torchvision import transforms
|
||
from taihuyuan_roof.dataloaders import custom_transforms as tr
|
||
|
||
import numpy as np
|
||
import random
|
||
import torch
|
||
|
||
seed = 7
|
||
random.seed(seed)
|
||
np.random.seed(seed)
|
||
torch.manual_seed(seed) # 为CPU设置随机种子
|
||
torch.cuda.manual_seed(seed) # 为当前GPU设置随机种子
|
||
|
||
class datasets_rooftaihuyuan(Dataset):
|
||
def __init__(self,
|
||
args,
|
||
base_dir=Path.db_root_dir("thy_roof"),
|
||
split='train',
|
||
isAug=False
|
||
):
|
||
super(datasets_rooftaihuyuan, self).__init__()
|
||
|
||
self.args = args
|
||
self.resize = args.resize
|
||
self.crop_size = args.crop_size # 是单数比如256，
|
||
self.flip_prob = args.flip_prob # 是0-1
|
||
|
||
self.isAug = isAug
|
||
self._base_dir = base_dir
|
||
self._image_dir = os.path.join(self._base_dir, 'images1')
|
||
self._cat_dir = os.path.join(self._base_dir, 'labels1')
|
||
|
||
if isinstance(split, str):
|
||
self.split = [split]
|
||
else:
|
||
split.sort()
|
||
self.split = split
|
||
|
||
_splits_dir = "/home/qiqq/q3dl/datalinan/taihuyuan_roof/"
|
||
|
||
self.im_ids = []
|
||
self.images = []
|
||
self.categories = []
|
||
|
||
for splt in self.split:
|
||
with open(os.path.join(os.path.join(_splits_dir, splt + '.txt')), "r") as f:
|
||
lines = f.read().splitlines()
|
||
|
||
for ii, line in enumerate(lines):
|
||
_image = os.path.join(self._image_dir, line + ".png") # 注意这个地方有的图片可能是jpg有的可能是png自己看着改
|
||
_cat = os.path.join(self._cat_dir, line + ".png") # 注意格式
|
||
# print(_image)
|
||
# print(_cat)
|
||
assert os.path.isfile(_image)
|
||
assert os.path.isfile(_cat)
|
||
# print("通过")
|
||
self.im_ids.append(line)
|
||
self.images.append(_image)
|
||
self.categories.append(_cat)
|
||
|
||
assert (len(self.images) == len(self.categories))
|
||
|
||
# Display stats
|
||
print('Number of images in {}: {:d}'.format(split, len(self.images)))
|
||
|
||
def __len__(self):
|
||
return len(self.images)
|
||
|
||
def __getitem__(self, index):
|
||
_img, _target = self._make_img_gt_point_pair(index)
|
||
sample = {'image': _img, 'label': _target}
|
||
|
||
for split in self.split:
|
||
if split == "trainbat":
|
||
if self.isAug: # 训练的时候使用数据增强验证的时候不需要
|
||
return self.AugTrain(sample)
|
||
else:
|
||
return self.transform_tr(sample)
|
||
elif split == 'valbat':
|
||
return self.transform_val(sample)
|
||
elif split == 'test1':
|
||
return self.transform_val(sample)
|
||
|
||
def _make_img_gt_point_pair(self, index):
|
||
_img = Image.open(self.images[index]).convert('RGB')
|
||
_target = Image.open(self.categories[index])
|
||
|
||
return _img, _target
|
||
|
||
|
||
def transform_tr(self, sample):
|
||
composed_transforms = transforms.Compose([
|
||
tr.Resize(self.args.resize), # 先缩放要不然原图太大了进不去
|
||
# tr.Normalize(mean=(0.231, 0.217, 0.22), std=(0.104, 0.086, 0.085)), # 针对whdld数据集的
|
||
tr.Normalize_simple(),
|
||
tr.ToTensor()])
|
||
|
||
return composed_transforms(sample)
|
||
|
||
|
||
def AugTrain(self, sample):
|
||
'''
|
||
这个版本的AugTrain使用的 Zhenchao Jin版本对numpy进行操作
|
||
:param sample:
|
||
:return:
|
||
|
||
随机裁剪，随机翻转，随机颜色增强，随机旋转这四个就够了以后别再浪费时间在这上边了
|
||
'''
|
||
composed_transforms = transforms.Compose([
|
||
tr.Tonumpy(),
|
||
# tr.RandomCrop(crop_size=self.crop_size),
|
||
tr.RandomFlip(),
|
||
tr.PhotoMetricDistortion(),
|
||
tr.RandomRotation(),
|
||
|
||
tr.Resize2(self.resize,scale_range=None,keep_ratio=False),
|
||
tr.Normalize_simple(),
|
||
tr.ToTensor()])
|
||
return composed_transforms(sample)
|
||
|
||
def transform_val(self, sample):
|
||
composed_transforms = transforms.Compose([
|
||
tr.Resize(self.args.resize),
|
||
# tr.Normalize(mean=(0.052, 0.05, 0.05), std=(0.027, 0.026, 0.026)),
|
||
tr.Normalize_simple(),
|
||
|
||
tr.ToTensor()])
|
||
return composed_transforms(sample)
|
||
|
||
|
||
def __str__(self):
|
||
return 'linan(split=' + str(self.split) + ')'
|
||
#
|
||
|
||
if __name__ == '__main__':
|
||
a =np.random.randn(3,3)
|
||
# print(a)
|
||
|
||
|
||
from taihuyuan_roof.dataloaders.utils import decode_segmap
|
||
from torch.utils.data import DataLoader
|
||
# import matplotlib.pyplot as plt
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser()
|
||
args = parser.parse_args()
|
||
args.resize = (512, 512)
|
||
args.crop_size = 480 #在验证的时候没有用 crop和filp 只是为了不报错
|
||
args.flip_prob = 0.5
|
||
|
||
voc_train = datasets_rooftaihuyuan(args, split='trainbat',isAug=True)
|
||
voc_val = datasets_rooftaihuyuan(args, split='valbat', isAug=False)
|
||
print(len(voc_train))
|
||
print(len(voc_val))
|
||
# for i in voc_train:
|
||
# print(type(i))
|
||
|
||
dataloader = DataLoader(voc_train, batch_size=1, shuffle=False, num_workers=0)
|
||
|
||
for ii, sample in enumerate(dataloader):
|
||
for jj in range(sample["image"].size()[0]):
|
||
img = sample['image'].numpy()
|
||
gt = sample['label'].numpy()
|
||
# gt=gt-1
|
||
# print(gt)
|
||
tmp = np.array(gt[jj]).astype(np.uint8)
|
||
segmap = decode_segmap(tmp, dataset='pascal_customer')
|
||
img_tmp = np.transpose(img[jj], axes=[1, 2, 0])
|
||
img_tmp *= 255.0
|
||
img_tmp = img_tmp.astype(np.uint8)
|
||
plt.figure()
|
||
plt.title('img')
|
||
plt.imshow(img_tmp)
|
||
plt.show()
|
||
plt.figure()
|
||
plt.title('label')
|
||
plt.imshow(segmap)
|
||
plt.show()
|
||
|
||
# if ii == 1:
|
||
# break
|
||
|