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)
img, mask = self.transforms(self.phase,
img,
mask,
self.mean,
self.std,
crop=self.crop,
height=self.height,
width=self.width)
mask = mask.permute(2, 0, 1)
return img, mask
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)
img, mask = self.transforms(self.phase,
img,
mask,
self.mean,
self.std,
crop=self.crop,
height=self.height,
width=self.width)
mask = mask.permute(2, 0, 1) # 4x256x1600
mask = mask.view(mask.size(0), -1)
mask = torch.sum(mask, dim=1)
mask = mask > 0
return img, mask.float()
if __name__ == "__main__":
data_folder = "datasets/Steel_data"
df_path = "datasets/Steel_data/train.csv"
mean = (0.485, 0.456, 0.406)
std = (0.229, 0.224, 0.225)
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)
file_names = df.index.tolist()
for index in range(len(file_names)):
image_id, mask = make_mask(index, df)
image_path = os.path.join(data_folder, 'train_images', image_id)
image = cv2.imread(image_path)
image_aug, mask_aug = data_augmentation(image,
mask,
crop=True,
height=256,
width=400)
normalize = Normalize(mean=mean, std=std)
image = normalize(image=image)['image']
image_aug = normalize(image=image_aug)['image']
image_mask = image_with_mask_numpy(deepcopy(image), mask, mean,
std)['image']
image_aug_mask = image_with_mask_numpy(deepcopy(image_aug), mask_aug,
mean, std)['image']