From 71edc600ead2d6ac2b33c6ac67a903ad724d7db7 Mon Sep 17 00:00:00 2001 From: XiangqianMa Date: Wed, 25 Sep 2019 19:40:05 +0800 Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E6=95=B0=E6=8D=AE=E9=9B=86?= =?UTF-8?q?=E5=8A=A0=E8=BD=BD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 4 +- datasets/__init__.py | 0 datasets/steel_dataset.py | 172 ++++++++++++++++++ utils/__init__.py | 0 ...clear-mask-visualization-and-simple-eda.py | 1 - utils/rle_parse.py | 33 ++++ 6 files changed, 207 insertions(+), 3 deletions(-) create mode 100644 datasets/__init__.py create mode 100644 datasets/steel_dataset.py create mode 100644 utils/__init__.py rename {datasets => utils}/clear-mask-visualization-and-simple-eda.py (99%) create mode 100644 utils/rle_parse.py diff --git a/README.md b/README.md index 03cedab..321f835 100644 --- a/README.md +++ b/README.md @@ -1,9 +1,9 @@ Code for Kaggle [Steel Defect Detection](https://www.kaggle.com/c/severstal-steel-defect-detection) ## Datasets Prepare -Download steel datasets from [here](https://www.kaggle.com/c/severstal-steel-defect-detection/data) , unzip and put them into `../input` directory. +Download steel datasets from [here](https://www.kaggle.com/c/severstal-steel-defect-detection/data) , unzip and put them into `../Input` directory. -Structure of the `../input` folder can be like: +Structure of the `../Input` folder can be like: ``` test_images diff --git a/datasets/__init__.py b/datasets/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/datasets/steel_dataset.py b/datasets/steel_dataset.py new file mode 100644 index 0000000..05c7efa --- /dev/null +++ b/datasets/steel_dataset.py @@ -0,0 +1,172 @@ +# coding: utf-8 + +import os +import cv2 +import pdb +import time +import warnings +import random +import numpy as np +import pandas as pd +from tqdm import tqdm +from torch.optim.lr_scheduler import ReduceLROnPlateau +from sklearn.model_selection import train_test_split, StratifiedKFold +import torch +import torch.nn as nn +from torch.nn import functional as F +import torch.optim as optim +import torch.backends.cudnn as cudnn +from torch.utils.data import DataLoader, Dataset, sampler +from matplotlib import pyplot as plt +from albumentations import (HorizontalFlip, ShiftScaleRotate, Normalize, Resize, Compose, GaussNoise) +from albumentations.torch import ToTensor +from utils.rle_parse import mask2rle, make_mask +warnings.filterwarnings("ignore") +seed = 69 +random.seed(seed) +os.environ["PYTHONHASHSEED"] = str(seed) +np.random.seed(seed) +torch.cuda.manual_seed(seed) +torch.backends.cudnn.deterministic = True + + +# Dataset +class SteelDataset(Dataset): + def __init__(self, df, data_folder, mean, std, phase): + self.df = df + self.root = data_folder + self.mean = mean + self.std = std + self.phase = phase + self.transforms = get_transforms(phase, mean, std) + self.fnames = self.df.index.tolist() + + def __getitem__(self, idx): + image_id, mask = make_mask(idx, self.df) + image_path = os.path.join(self.root, "train_images", image_id) + img = cv2.imread(image_path) + augmented = self.transforms(image=img, mask=mask) + img = augmented['image'] + mask = augmented['mask'] # 1x256x1600x4 + mask = mask[0].permute(2, 0, 1) # 1x4x256x1600 + return img, mask + + def __len__(self): + return len(self.fnames) + + +def get_transforms(phase, mean, std): + list_transforms = [] + if phase == "train": + list_transforms.extend( + [ + HorizontalFlip(p=0.5), # only horizontal flip as of now + ] + ) + list_transforms.extend( + [ + Normalize(mean=mean, std=std, p=1), + ToTensor(), + ] + ) + list_trfms = Compose(list_transforms) + return list_trfms + + +def provider( + data_folder, + df_path, + mean=None, + std=None, + batch_size=8, + num_workers=4, + n_splits=0, +): + """返回数据加载器 + + Args: + data_folder: 数据集根目录 + df_path: csv文件路径 + mean: 数据集各通道均值 + std: 数据集各通道标准差 + batch_size + num_workers + n_split: 交叉验证折数，为1时不使用交叉验证 + + Return: + dataloadrs: list，该list中的每一个元素为list，元素list中保存训练集和验证集 + """ + df = pd.read_csv(df_path) + # https://www.kaggle.com/amanooo/defect-detection-starter-u-net + df['ImageId'], df['ClassId'] = zip(*df['ImageId_ClassId'].str.split('_')) + df['ClassId'] = df['ClassId'].astype(int) + df = df.pivot(index='ImageId',columns='ClassId',values='EncodedPixels') + df['defects'] = df.count(axis=1) + + # 将数据集划分为n_split份 + train_dfs = list() + val_dfs = list() + if n_splits != 1: + skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=69) + # 按照样本包含的损伤的类别的数目进行分层 + for train_df_index, val_df_index in skf.split(df, df['defects']): + train_dfs.append(df.loc[df.index[train_df_index]]) + val_dfs.append(df.loc[df.index[val_df_index]]) + else: + df_temp = train_test_split(df, test_size=0.2, stratify=df["defects"], random_state=69) + train_dfs.append(df_temp[0]) + val_dfs.append(df_temp[1]) + # 生成dataloader + dataloaders = list() + for df_index, (train_df, val_df) in enumerate(zip(train_dfs, val_dfs)): + train_dataset = SteelDataset(train_df, data_folder, mean, std, 'train') + val_dataset = SteelDataset(val_df, data_folder, mean, std, 'val') + train_dataloader = DataLoader( + train_dataset, + batch_size=batch_size, + num_workers=num_workers, + pin_memory=True, + shuffle=True + ) + val_dataloader = DataLoader( + val_dataset, + batch_size=batch_size, + num_workers=num_workers, + pin_memory=True, + shuffle=True + ) + dataloaders.append([train_dataloader, val_dataloader]) + + return dataloaders + + +if __name__ == "__main__": + data_folder = './Steel_data' + df_path = '/home/apple/program/MXQ/Competition/Kaggle/Steal-Defect/Kaggle-Steel-Defect-Detection/datasets/Steel_data/train.csv' + mean=(0.485, 0.456, 0.406) + std=(0.229, 0.224, 0.225) + batch_size = 8 + num_workers = 4 + n_splits = 1 + dataloader = provider(data_folder, df_path, mean, std, batch_size, num_workers, n_splits) + for fold_index, [train_dataloader, val_dataloader] in enumerate(dataloader): + train_bar = tqdm(train_dataloader) + class_color = [[255, 0, 0], [0, 255, 0], [0, 0, 255], [139, 0, 139]] + for images, targets in train_bar: + image = images[0] + for i in range(3): + image[i] = image[i] * std[i] + image[i] = image[i] + mean[i] + + target = targets[0] + for i in range(target.size(0)): + target_0 = target[i] * class_color[i][0] + target_1 = target[i] * class_color[i][1] + target_2 = target[i] * class_color[i][2] + mask = torch.stack([target_0, target_1, target_2], dim=0) + image += mask + image = image.permute(1, 2, 0).numpy() + cv2.imshow('win', image) + cv2.waitKey(0) + + pass diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/datasets/clear-mask-visualization-and-simple-eda.py b/utils/clear-mask-visualization-and-simple-eda.py similarity index 99% rename from datasets/clear-mask-visualization-and-simple-eda.py rename to utils/clear-mask-visualization-and-simple-eda.py index 650ff3d..4b76451 100644 --- a/datasets/clear-mask-visualization-and-simple-eda.py +++ b/utils/clear-mask-visualization-and-simple-eda.py @@ -19,7 +19,6 @@ import matplotlib.pyplot as plt # get_ipython().run_line_magic('matplotlib', 'inline') # plt.rcParams["font.size"] = 15 -import seaborn as sns from collections import Counter from PIL import Image import math diff --git a/utils/rle_parse.py b/utils/rle_parse.py new file mode 100644 index 0000000..b268557 --- /dev/null +++ b/utils/rle_parse.py @@ -0,0 +1,33 @@ +import numpy as np +import pandas as pd + + +# https://www.kaggle.com/paulorzp/rle-functions-run-lenght-encode-decode +def mask2rle(img): + ''' + img: numpy array, 1 -> mask, 0 -> background + Returns run length as string formated + ''' + pixels= img.T.flatten() + pixels = np.concatenate([[0], pixels, [0]]) + runs = np.where(pixels[1:] != pixels[:-1])[0] + 1 + runs[1::2] -= runs[::2] + return ' '.join(str(x) for x in runs) + +def make_mask(row_id, df): + '''Given a row index, return image_id and mask (256, 1600, 4) from the dataframe `df`''' + fname = df.iloc[row_id].name + labels = df.iloc[row_id][:4] + masks = np.zeros((256, 1600, 4), dtype=np.float32) # float32 is V.Imp + # 4:class 1～4 (ch:0～3) + + for idx, label in enumerate(labels.values): + if label is not np.nan: + label = label.split(" ") + positions = map(int, label[0::2]) + length = map(int, label[1::2]) + mask = np.zeros(256 * 1600, dtype=np.uint8) + for pos, le in zip(positions, length): + mask[pos:(pos + le)] = 1 + masks[:, :, idx] = mask.reshape(256, 1600, order='F') + return fname, masks \ No newline at end of file