diff --git a/new_unet/net.py b/new_unet/net.py new file mode 100644 index 0000000000..336633f93e --- /dev/null +++ b/new_unet/net.py @@ -0,0 +1,97 @@ +import torch +from torch import nn +from torch.nn import functional as F + + +# 卷积类 +class Conv_Block(nn.Module): + + def __init__(self, in_channels, out_channels): + super(Conv_Block, self).__init__() + self.layer = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=3,stride=1, padding=1,padding_mode='reflect',bias=False), + nn.BatchNorm2d(out_channels), + nn.Dropout(0.3), + nn.LeakyReLU(), + # 第一个卷积 + nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, padding_mode='reflect', + bias=False), + nn.BatchNorm2d(out_channels), + nn.Dropout(0.3), + nn.LeakyReLU(), + # 第二个卷积同理 + ) + def forward(self, x): + return self.layer(x) + + +#下采样类 +class DownSample(nn.Module): + def __init__(self, channel): + super(DownSample, self).__init__() + # 序列构造器 + self.layer = nn.Sequential( + nn.Conv2d(channel, channel, kernel_size=3, stride=2, padding=1, padding_mode='reflect', bias=False), + nn.BatchNorm2d(channel),#可以用也可以不用 + nn.LeakyReLU() + ) + + def forward(self, x): + return self.layer(x) + + +#上采样 得到特征图，和同层的下采样图进行拼接，然后再进行卷积 +class UpSample(nn.Module): + def __init__(self,channel): + super(UpSample, self).__init__() + #转置卷积与插值法 + #用1+1的卷积降通道 + self.layer = nn.Conv2d(channel, channel//2, kernel_size=1, stride=1) + + def forward(self, x,feature_map): #feature_map 是同层之前下采样的特征图 + up = F.interpolate(x,scale_factor=2,mode='nearest')#scale_factor=2 变成原来的2倍 插值法 + out = self.layer(up) + return torch.cat([out,feature_map],dim=1) #n c h w 0 1 2 3 + + +class UNet(nn.Module): + def __init__(self): + super(UNet, self).__init__() + self.c1=Conv_Block(3, 64) + self.d1=DownSample(64) + self.c2=Conv_Block(64, 128) + self.d2=DownSample(128) + self.c3=Conv_Block(128, 256) + self.d3=DownSample(256) + self.c4=Conv_Block(256, 512) + self.d4=DownSample(512) + self.c5=Conv_Block(512, 1024) + self.u1=UpSample(1024) + self.c6=Conv_Block(1024, 512) + self.u2=UpSample(512) + self.c7=Conv_Block(512, 256) + self.u3=UpSample(256) + self.c8=Conv_Block(256, 128) + self.u4=UpSample(128) + self.c9=Conv_Block(128, 64) + self.out=nn.Conv2d(64, 3,kernel_size=1, stride=1, padding=1) + self.Th=nn.Sigmoid()#激活函数，虽然是彩色图像，只需要对图像进行二分类即可，也就是讲一个像素点分为黑色或白色 + + def forward(self, x): + R1=self.c1(x) + R2=self.c2(self.d1(R1))#先对R1进行下采样，然后再进行第二次的一个卷积 + R3=self.c3(self.d2(R2)) + R4=self.c4(self.d3(R3)) + R5=self.c5(self.d4(R4)) + O1=self.c6(self.u1(R5,R4)) #先将R5上采样，并且和R4进行拼接，然后再进行卷积 + O2=self.c7(self.u2(O1,R3)) + O3=self.c8(self.u3(O2,R2)) + O4=self.c9(self.u4(O3,R1)) + + return self.Th(self.out(O4)) + +#测试，输出模型的形状，看看和我们设计的样子是否一样 +if __name__ == '__main__': + x = torch.randn(2, 3, 256, 256) + net = UNet() + print(net(x).shape) \ No newline at end of file diff --git a/scripts/download_data.bat b/scripts/download_data.bat new file mode 100644 index 0000000000..672729cba6 --- /dev/null +++ b/scripts/download_data.bat @@ -0,0 +1,28 @@ +@echo off +setlocal enabledelayedexpansion + +if not exist "%userprofile%\.kaggle\kaggle.json" ( + set /p USERNAME=Kaggle username: + echo. + set /p APIKEY=Kaggle API key: + + mkdir "%userprofile%\.kaggle" + echo {"username":"!USERNAME!","key":"!APIKEY!"} > "%userprofile%\.kaggle\kaggle.json" + attrib +R "%userprofile%\.kaggle\kaggle.json" +) + +pip install kaggle --upgrade + +kaggle competitions download -c carvana-image-masking-challenge -f train_hq.zip +powershell Expand-Archive train_hq.zip -DestinationPath data\imgs +move data\imgs\train_hq\* data\imgs\ +rmdir /s /q data\imgs\train_hq +del /q train_hq.zip + +kaggle competitions download -c carvana-image-masking-challenge -f train_masks.zip +powershell Expand-Archive train_masks.zip -DestinationPath data\masks +move data\masks\train_masks\* data\masks\ +rmdir /s /q data\masks\train_masks +del /q train_masks.zip + +exit /b 0 diff --git a/train.py b/train.py index 6c5a51a29c..17b3026195 100644 --- a/train.py +++ b/train.py @@ -19,8 +19,11 @@ from utils.data_loading import BasicDataset, CarvanaDataset from utils.dice_score import dice_loss -dir_img = Path('./data/imgs/') -dir_mask = Path('./data/masks/') +# dir_img = Path('./data/imgs/') +# dir_mask = Path('./data/masks/') +#注释掉原来的地址 +dir_img = Path('../Dataset/TrainDataset/Image/') +dir_mask = Path('../Dataset/TrainDataset/GT/') dir_checkpoint = Path('./checkpoints/') @@ -33,7 +36,7 @@ def train_model( val_percent: float = 0.1, save_checkpoint: bool = True, img_scale: float = 0.5, - amp: bool = False, + amp: bool = True, weight_decay: float = 1e-8, momentum: float = 0.999, gradient_clipping: float = 1.0, diff --git a/utils/data_loading.py b/utils/data_loading.py index 11296e78bb..3a69a2c4ca 100644 --- a/utils/data_loading.py +++ b/utils/data_loading.py @@ -12,9 +12,20 @@ from torch.utils.data import Dataset from tqdm import tqdm +def keep_image_size_open(path, size=(256, 256)): + img = Image.open(path) + temp = max(img.size) + mask = Image.new('RGB',(temp, temp), (0, 0, 0)) + mask.paste(img, (0, 0)) + mask = mask.resize(size) + return mask + def load_image(filename): ext = splitext(filename)[1] + + # splitext(filename)是os.path模块中的函数，用于将文件名分离成文件名和扩展名两个部分。 + # splitext(filename)[1]提取扩展名部分，比如.jpg、.png、.npy等。这里的ext变量将保存这个扩展名。 if ext == '.npy': return Image.fromarray(np.load(filename)) elif ext in ['.pt', '.pth']: @@ -24,14 +35,47 @@ def load_image(filename): def unique_mask_values(idx, mask_dir, mask_suffix): - mask_file = list(mask_dir.glob(idx + mask_suffix + '.*'))[0] + mask_file = list(mask_dir.glob(idx + mask_suffix + '.png'))[0] + # print(f"找到遮罩文件: {mask_file}") + # 得到mask_dir目录下 idx.png 格式的名字 + #----------------修改 开始------------ + # try: + # # 获取第一个匹配文件 + # print(f"找到遮罩文件: {mask_file}") + # except IndexError: + # print("没有找到与模式匹配的遮罩文件。") + # ----------------修改 结束------------ mask = np.asarray(load_image(mask_file)) + a = np.unique(mask) + # print(f'idx: {idx} a {a}') + # if a == [ 0 255] : + # print(f'idx: {idx} a {a}') + # else: + # print(f'-------------------idx: {idx} a {a}--------------------') + + # print(f"{mask_file} np.unique(mask): {np.unique(mask)}") + # 得到像素中的不同数 if mask.ndim == 2: + # print("ndim == 2") cod10k是这个 + + # a = np.unique(mask) + # print(f"{mask_file} a.ndim: {a.ndim} ") + # print(f"{mask_file} np.unique(mask): {np.unique(mask)}") + # if a.ndim == 2: + # print(f'a.ndim=2*****************') + # # print(f"{mask_file} np.unique(mask): {np.unique(mask)}") + # if a.ndim>2 : + # print(f'a.ndim>2*****************') + # print(f"{mask_file} np.unique(mask): {np.unique(mask)}") return np.unique(mask) + # return mask elif mask.ndim == 3: mask = mask.reshape(-1, mask.shape[-1]) + print(f"{mask_file} np.unique(mask): {np.unique(mask)}") + # print("ndim == 3") return np.unique(mask, axis=0) else: + print(f"{mask_file} np.unique(mask): {np.unique(mask)}") raise ValueError(f'Loaded masks should have 2 or 3 dimensions, found {mask.ndim}') @@ -44,6 +88,7 @@ def __init__(self, images_dir: str, mask_dir: str, scale: float = 1.0, mask_suff self.mask_suffix = mask_suffix self.ids = [splitext(file)[0] for file in listdir(images_dir) if isfile(join(images_dir, file)) and not file.startswith('.')] + # self.ids 是照片的名称，不带文件格式的文件名。 if not self.ids: raise RuntimeError(f'No input file found in {images_dir}, make sure you put your images there') @@ -55,7 +100,45 @@ def __init__(self, images_dir: str, mask_dir: str, scale: float = 1.0, mask_suff total=len(self.ids) )) + + # print(f'unique[0] {unique[0]} ') + # + # print(f'self.ids[0] {self.ids[0]}') + # + # mask_file = list(mask_dir.glob(self.ids[0] + mask_suffix + '.*'))[0] + # print(f'mask_file {mask_file}') + # mask = np.asarray(load_image(mask_file)) + # print(f'mask {mask}') + # np_mask = np.unique(mask) + # print(f'np_mask {np_mask}') + # aa=unique_mask_values(self.ids[0], self.mask_dir, self.mask_suffix) + # print(f'unique_mask_values aa {aa}') + # tqdm(unique_mask_values(self.ids[0], self.mask_dir, self.mask_suffix),total=1) + # unique1 = list( (tqdm(unique_mask_values(self.ids[0], self.mask_dir, self.mask_suffix),total=1) )) + # print(f'unique1 {unique1}') + # unique2 = list(unique_mask_values(self.ids[0], self.mask_dir, self.mask_suffix)) + # print(f'unique1 {unique2}') + # + print(f'unique[0] {unique[0]} ') + print(f'unique[1] {unique[1]} ') + # unique_mask_values(self.ids[1], self.mask_dir, self.mask_dir) + print(f'unique[2] {unique[2]} ') + print(f'unique[3] {unique[3]} ') + # unique_mask_values(self.ids[2], self.mask_dir, self.mask_dir) + print(f'unique[4] {unique[4]} ') + print(f'unique[5] {unique[5]} ') + print(f'unique[6] {unique[6]} ') + print(f'unique[7] {unique[7]} ') + print(f'unique[8] {unique[8]} ') + # unique_mask_values(self.ids[6], self.mask_dir, self.mask_dir) + c = np.concatenate(unique) + # a = np.unique(np.concatenate(unique), axis=0).tolist() + # self.mask_values = list(sorted(a)) + + self.mask_values = list(sorted(np.unique(np.concatenate(unique), axis=0).tolist())) + # self.mask_values = [0,255] + logging.info(f'Unique mask values: {self.mask_values}') def __len__(self): @@ -63,6 +146,7 @@ def __len__(self): @staticmethod def preprocess(mask_values, pil_img, scale, is_mask): + # w, h = pil_img.size newW, newH = int(scale * w), int(scale * h) assert newW > 0 and newH > 0, 'Scale is too small, resized images would have no pixel' @@ -70,6 +154,7 @@ def preprocess(mask_values, pil_img, scale, is_mask): img = np.asarray(pil_img) if is_mask: + #如果是mask图片，就生成一张大小与图像相同的空遮挡 mask = np.zeros((newH, newW), dtype=np.int64) for i, v in enumerate(mask_values): if img.ndim == 2: @@ -92,12 +177,19 @@ def preprocess(mask_values, pil_img, scale, is_mask): def __getitem__(self, idx): name = self.ids[idx] - mask_file = list(self.mask_dir.glob(name + self.mask_suffix + '.*')) - img_file = list(self.images_dir.glob(name + '.*')) + # mask_file = list(self.mask_dir.glob(name + self.mask_suffix + '.*')) + # img_file = list(self.images_dir.glob(name + '.*')) + #将数据集图片的格式改成自己的，这里都是.jpg + #images_dir - dir_img - '../Dataset/TrainDataset/Image/' + #mask_file - dir_mask - '../Dataset/TrainDataset/GT/' + mask_file = list(self.mask_dir.glob(name + self.mask_suffix + '.png')) + #将文件路径变成list存储 + img_file = list(self.images_dir.glob(name + '.jpg')) + # assert len(img_file) == 1, f'Either no image or multiple images found for the ID {name}: {img_file}' assert len(mask_file) == 1, f'Either no mask or multiple masks found for the ID {name}: {mask_file}' - mask = load_image(mask_file[0]) + mask = load_image(mask_file[0])#取第一个文件路径并加载 img = load_image(img_file[0]) assert img.size == mask.size, \ @@ -114,4 +206,19 @@ def __getitem__(self, idx): class CarvanaDataset(BasicDataset): def __init__(self, images_dir, mask_dir, scale=1): - super().__init__(images_dir, mask_dir, scale, mask_suffix='_mask') + super().__init__(images_dir, mask_dir, scale, mask_suffix='') + +def unique_values(idx, mask_dir): + # imgs = os.listdir(root) + imgs = list(mask_dir.glob(idx + '.png'))[0] + concat_unique = np.empty(1) + for imgpath in imgs: + img = np.asarray(Image.open(imgs)) + # 得到像素中的不同数 + unique = np.unique(img) + # 对其进行拼接 + concat_unique = np.concatenate([concat_unique, unique]) + # 对拼接后的图片进行再次求不同像素，即全部文件中不同像素数，排序后返回 + return list(sorted(np.unique(concat_unique))) +# if __name__ == '__main__': +