I am trying to use this script for my own use in a dataset from Kaggle. I have created the masks images already and do not need some of the code shown in the page. However, I am getting an error while trying to run my script which I think is a simple issue. It is shown below. Please find the sample data here.
Would anyone be able to help me in this matter please.
#import libraries and define paths
import os
import sys
import json
import numpy as np
from skimage.io import imread, imsave
from skimage.transform import resize
from skimage.draw import draw
import matplotlib.pyplot as plt
import albumentations as A
import cv2
import random
import torch
from albumentations.pytorch import ToTensorV2
from unet_network import UNet
import time
import copy
img_longest_size = 2308
data_dir = '/Users/schrotermichael/Downloads/hubmap-organ-segmentation/'
#image_dir = os.path.join(data_dir, "images_pike")
image_small_dir = os.path.join(data_dir, "train_img_cropped")
#mask_dir = os.path.join(data_dir, "images_pike_masks")
mask_small_dir = os.path.join(data_dir, "train_anno_cropped")
#via_proj = os.path.join(data_dir, "via_project.json")
#define dataset for pytorch
class PikeDataset(torch.utils.data.Dataset):
def __init__(self, images_directory, masks_directory, mask_filenames, transform=None):
self.images_directory = images_directory
self.masks_directory = masks_directory
self.mask_filenames = mask_filenames
self.transform = transform
def __len__(self):
return len(self.mask_filenames)
def __getitem__(self, idx):
mask_filename = self.mask_filenames[idx]
image = cv2.imread(os.path.join(self.images_directory, mask_filename))
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
mask = cv2.imread(os.path.join(self.masks_directory, mask_filename),
cv2.IMREAD_UNCHANGED)
mask = np.expand_dims(mask, 2)
mask = mask.astype(np.float32)
if self.transform is not None:
transformed = self.transform(image=image, mask=mask)
image = transformed["image"]
mask = transformed["mask"]
return image, mask
#define augmentations
train_transform = A.Compose([
A.Resize(400, 400, always_apply=True),
A.RandomCrop(height=200, width = 200, p=0.2),
A.PadIfNeeded(min_height=400, min_width=400, border_mode=cv2.BORDER_CONSTANT,
always_apply=True),
A.VerticalFlip(p=0.2),
A.Blur(p=0.2),
A.RandomRotate90(p=0.2),
A.ShiftScaleRotate(p=0.2, border_mode=cv2.BORDER_CONSTANT),
A.RandomBrightnessContrast(p=0.2),
A.RandomSunFlare(p=0.2, src_radius=200),
A.RandomShadow(p=0.2),
A.RandomFog(p=0.2),
A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
ToTensorV2(transpose_mask=True)
]
)
val_transform = A.Compose([
A.Resize(400, 400, always_apply=True),
A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
ToTensorV2(transpose_mask=True)
]
)
#split files at random into training and validation datasets
mask_filenames = list(sorted(os.listdir(mask_small_dir)))
random.seed(43)
random.shuffle(mask_filenames)
n_val = 5
val_mask_filenames = mask_filenames[:5]
train_mask_filenames = mask_filenames[5:]
train_dataset = PikeDataset(image_small_dir, mask_small_dir, train_mask_filenames,
transform=train_transform)
val_dataset = PikeDataset(image_small_dir, mask_small_dir, val_mask_filenames,
transform=val_transform)
#dice loss function
class DiceLoss(torch.nn.Module):
def __init__(self):
super(DiceLoss, self).__init__()
self.smooth = 1.0
def forward(self, y_pred, y_true):
assert y_pred.size() == y_true.size()
y_pred = y_pred[:, 0].contiguous().view(-1)
y_true = y_true[:, 0].contiguous().view(-1)
intersection = (y_pred * y_true).sum()
dsc = (2.*intersection+ elf.smooth)/(y_pred.sum()+y_true.sum()+self.smooth)
return 1. - dsc
#function to return data loaders
def data_loaders(dataset_train, dataset_valid, bs = 5, workers = 0):
def worker_init(worker_id):
np.random.seed(42 + worker_id)
loader_train = torch.utils.data.DataLoader(
dataset_train,
batch_size=bs,
shuffle=True,
drop_last=True,
num_workers=workers,
worker_init_fn=worker_init,
)
loader_valid = torch.utils.data.DataLoader(
dataset_valid,
batch_size=bs,
drop_last=False,
shuffle=False,
num_workers=workers,
worker_init_fn=worker_init,
)
return {"train": loader_train, "valid": loader_valid}
#define datasets and load network from file
loaders = data_loaders(train_dataset, val_dataset)
device = torch.device("cpu" if not torch.cuda.is_available() else "cuda")
unet = UNet(in_channels=3, out_channels=1, init_features=8)
unet.to(device)
dsc_loss = DiceLoss()
best_validation_dsc = 1.0
lr = 0.001
epochs = 200
optimizer = torch.optim.Adam(unet.parameters(), lr=lr)
best_val_true = []
best_val_pred = []
epoch_loss = {"train": [], "valid": []}
#The below code block gave me an error.
#training loop
for epoch in range(epochs):
print('-' * 100)
since = time.time()
best_model_wts = copy.deepcopy(unet.state_dict())
for phase in ["train", "valid"]:
if phase == "train":
unet.train()
else:
unet.eval()
epoch_samples = 0
running_loss = 0
for i, data in enumerate(loaders[phase]):
x, y_true = data
x, y_true = x.to(device), y_true.to(device)
epoch_samples += x.size(0)
optimizer.zero_grad()
with torch.set_grad_enabled(phase == "train"):
y_pred = unet(x)
loss = dsc_loss(y_pred, y_true)
running_loss += loss.item()
if phase == "train":
loss.backward()
optimizer.step()
epoch_phase_loss = running_loss/epoch_samples
epoch_loss[phase].append(epoch_phase_loss)
if phase == "valid" and epoch_phase_loss < best_validation_dsc:
best_validation_dsc = epoch_phase_loss
print("Saving best model")
best_model_wts = copy.deepcopy(unet.state_dict())
best_val_true = y_true.detach().cpu().numpy()
best_val_pred = y_pred.detach().cpu().numpy()
time_elapsed = time.time() - since
print("Epoch {}/{}: Train loss = {:4f} --- Valid loss = {:4f} --- Time: {:.0f}m {:.0f}s".format(epoch + 1, epochs, epoch_loss["train"][epoch], epoch_loss["valid"][epoch], time_elapsed // 60, time_elapsed % 60))
print("Best validation loss: {:4f}".format(best_validation_dsc))
torch.save(best_model_wts, os.path.join(data_dir, "unet.pt"))
Output of the above code block
----------------------------------------------------------------------------------------------------
---------------------------------------------------------------------------
error Traceback (most recent call last)
Input In [22], in <cell line: 2>()
14 epoch_samples = 0
15 running_loss = 0
---> 17 for i, data in enumerate(loaders[phase]):
18 x, y_true = data
19 x, y_true = x.to(device), y_true.to(device)
File ~/newEnv/lib/python3.10/site-packages/torch/utils/data/dataloader.py:681, in _BaseDataLoaderIter.__next__(self)
678 if self._sampler_iter is None:
679 # TODO(https://github.com/pytorch/pytorch/issues/76750)
680 self._reset() # type: ignore[call-arg]
--> 681 data = self._next_data()
682 self._num_yielded += 1
683 if self._dataset_kind == _DatasetKind.Iterable and \
684 self._IterableDataset_len_called is not None and \
685 self._num_yielded > self._IterableDataset_len_called:
File ~/newEnv/lib/python3.10/site-packages/torch/utils/data/dataloader.py:721, in _SingleProcessDataLoaderIter._next_data(self)
719 def _next_data(self):
720 index = self._next_index() # may raise StopIteration
--> 721 data = self._dataset_fetcher.fetch(index) # may raise StopIteration
722 if self._pin_memory:
723 data = _utils.pin_memory.pin_memory(data, self._pin_memory_device)
File ~/newEnv/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py:49, in _MapDatasetFetcher.fetch(self, possibly_batched_index)
47 def fetch(self, possibly_batched_index):
48 if self.auto_collation:
---> 49 data = [self.dataset[idx] for idx in possibly_batched_index]
50 else:
51 data = self.dataset[possibly_batched_index]
File ~/newEnv/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py:49, in <listcomp>(.0)
47 def fetch(self, possibly_batched_index):
48 if self.auto_collation:
---> 49 data = [self.dataset[idx] for idx in possibly_batched_index]
50 else:
51 data = self.dataset[possibly_batched_index]
Input In [18], in PikeDataset.__getitem__(self, idx)
20 mask = mask.astype(np.float32)
22 if self.transform is not None:
---> 23 transformed = self.transform(image=image, mask=mask)
24 image = transformed["image"]
25 mask = transformed["mask"]
File ~/newEnv/lib/python3.10/site-packages/albumentations/core/composition.py:205, in Compose.__call__(self, force_apply, *args, **data)
202 p.preprocess(data)
204 for idx, t in enumerate(transforms):
--> 205 data = t(force_apply=force_apply, **data)
207 if check_each_transform:
208 data = self._check_data_post_transform(data)
File ~/newEnv/lib/python3.10/site-packages/albumentations/core/transforms_interface.py:98, in BasicTransform.__call__(self, force_apply, *args, **kwargs)
93 warn(
94 self.get_class_fullname() + " could work incorrectly in ReplayMode for other input data"
95 " because its' params depend on targets."
96 )
97 kwargs[self.save_key][id(self)] = deepcopy(params)
---> 98 return self.apply_with_params(params, **kwargs)
100 return kwargs
File ~/newEnv/lib/python3.10/site-packages/albumentations/core/transforms_interface.py:111, in BasicTransform.apply_with_params(self, params, **kwargs)
109 target_function = self._get_target_function(key)
110 target_dependencies = {k: kwargs[k] for k in self.target_dependence.get(key, [])}
--> 111 res[key] = target_function(arg, **dict(params, **target_dependencies))
112 else:
113 res[key] = None
File ~/newEnv/lib/python3.10/site-packages/albumentations/core/transforms_interface.py:240, in DualTransform.apply_to_mask(self, img, **params)
239 def apply_to_mask(self, img: np.ndarray, **params) -> np.ndarray:
--> 240 return self.apply(img, **{k: cv2.INTER_NEAREST if k == "interpolation" else v for k, v in params.items()})
File ~/newEnv/lib/python3.10/site-packages/albumentations/augmentations/geometric/resize.py:179, in Resize.apply(self, img, interpolation, **params)
178 def apply(self, img, interpolation=cv2.INTER_LINEAR, **params):
--> 179 return F.resize(img, height=self.height, width=self.width, interpolation=interpolation)
File ~/newEnv/lib/python3.10/site-packages/albumentations/augmentations/functional.py:172, in preserve_channel_dim.<locals>.wrapped_function(img, *args, **kwargs)
169 @wraps(func)
170 def wrapped_function(img, *args, **kwargs):
171 shape = img.shape
--> 172 result = func(img, *args, **kwargs)
173 if len(shape) == 3 and shape[-1] == 1 and len(result.shape) == 2:
174 result = np.expand_dims(result, axis=-1)
File ~/newEnv/lib/python3.10/site-packages/albumentations/augmentations/geometric/functional.py:364, in resize(img, height, width, interpolation)
362 return img
363 resize_fn = _maybe_process_in_chunks(cv2.resize, dsize=(width, height), interpolation=interpolation)
--> 364 return resize_fn(img)
File ~/newEnv/lib/python3.10/site-packages/albumentations/augmentations/functional.py:297, in _maybe_process_in_chunks.<locals>.__process_fn(img)
295 img = np.dstack(chunks)
296 else:
--> 297 img = process_fn(img, **kwargs)
298 return img
error: OpenCV(4.6.0) /Users/runner/work/opencv-python/opencv-python/opencv/modules/imgproc/src/resize.cpp:3689: error: (-215:Assertion failed) !dsize.empty() in function 'resize'
The rest of the code:
#load weights to network
weights_path = data_dir + "unet.pt"
device = "cpu"
unet = UNet(in_channels=3, out_channels=1, init_features=8)
unet.to(device)
unet.load_state_dict(torch.load(weights_path, map_location=device))
#define augmentations
inference_transform = A.Compose([
A.Resize(400, 400, always_apply=True),
A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
ToTensorV2()
])
#define function for predictions
def predict(model, img, device):
model.eval()
with torch.no_grad():
images = img.to(device)
output = model(images)
predicted_masks = (output.squeeze() >= 0.5).float().cpu().numpy()
return(predicted_masks)
#define function to load image and output mask
def get_mask(img_path):
image = cv2.imread(img_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
original_height, original_width = tuple(image.shape[:2])
image_trans = inference_transform(image = image)
image_trans = image_trans["image"]
image_trans = image_trans.unsqueeze(0)
image_mask = predict(unet, image_trans, device)
image_mask = F.resize(image_mask, original_height, original_width,
interpolation=cv2.INTER_NEAREST)
return(image_mask)
#image example
example_path = "10044.jpeg"
image = cv2.imread(example_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
mask = get_mask(example_path)
masked_img = image*np.expand_dims(mask, 2).astype("uint8")
#plot the image, mask and multiplied together
fig, (ax1, ax2, ax3) = plt.subplots(3)
ax1.imshow(image)
ax2.imshow(mask)
ax3.imshow(masked_img)