Fastai: Error Message in learn.predict when using custom Datablock/ Dataloader - image-segmentation

Question

I need some help with my Fastai pipeline.
I want to do semantic segmentation on a 2 channel input image with augmentation.
I adapted my procedure from the good introduction in medium
I have 2 channel images that are saved as NumPy arrays (.npy) of the size 2x 426 x 476.
See my code below:
%matplotlib inline
import torch
print(torch.__version__)
print(torch.cuda.is_available())
import fastai
print(fastai.__version__)
# other imports
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
from PIL import Image
from fastai.vision.all import *
from sklearn.model_selection import StratifiedKFold
import torch
import fastai
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
from PIL import Image
from fastai.vision.all import *
imgs_path = Path('./data')
lbls_path = Path('./label1')
def open_im(fn, chnls=None, cls=torch.Tensor):
im = (np.load(fn)).astype('float32')
return cls(im)
# The map_filename function makes it easier to map from one folder to another by replacing strings
def map_filename(base_fn, str1, str2):
return Path(str(base_fn).replace(str1, str2))
# get items from both datasets
items = get_files('./data', extensions='.npy')
masks = get_files('./label1', extensions='.npy')
items_mask = masks
items = items
items
idx=2
img_pipe = Pipeline(open_im)
img = img_pipe(items[idx])
mask_pipe = Pipeline([partial(map_filename, str1='data', str2='label1'),
partial(open_im, cls=TensorMask)])
mask = mask_pipe(items_mask[idx])
print(img.shape, mask.shape)
_, ax = plt.subplots(1, 2, figsize=(12,5))
ax[0].imshow(img.permute(1, 2, 0)[..., :1]/20000)
mask.show(ctx=ax[1])
plt.show()
def show_img(tensor_img, ctx=None):
ctx = plt.subplot() if ctx is None else ctx
#normalize to fit between 0 and 1
if tensor_img.max() > 0:
tensor_img = tensor_img / tensor_img.max()
ctx.imshow(tensor_img.permute(1, 2, 0)[..., :1])
# To create this DataBlock we don't need to specify the get_items function
# because we will pass the list of files as the source
db = DataBlock(blocks=(TransformBlock([open_im, lambda x: x/10000]),
TransformBlock([partial(map_filename, str1='data', str2='label1'),
partial(open_im, cls=TensorMask)])),
splitter=RandomSplitter(valid_pct=0.2, seed=0)
)
db.summary(source=items)
ds = db.datasets(source=items)
dl = db.dataloaders(source=items, bs=1)
batch = dl.one_batch()
print(batch[0].shape, batch[1].shape)
import albumentations as A
import pdb
class SegmentationAlbumentationsTransform(ItemTransform):
# split_idx=0
def __init__(self, aug, **kwargs):
super().__init__(**kwargs)
self.aug = aug
def encodes(self, x):
img,mask = x
img = img/img.max()
aug = self.aug(image=np.array(img.permute(1,2,0)), mask=np.array(mask))
return TensorImage(aug['image'].transpose(2,0,1)), TensorMask(aug['mask'])
aug_pipe = A.Compose([
A.ShiftScaleRotate(p=.9),
A.HorizontalFlip(),
A.RandomCrop(384, 384),
A.Rotate(limit=(-90, 90)),
# A.RandomBrightnessContrast(contrast_limit=0.0, p=1., brightness_by_max=False)
])
# Create our class with this aug_pipe
aug = SegmentationAlbumentationsTransform(aug_pipe)
# And check the results
idx = 5
aug_number = 4
# Display original and some augmented samples
_, ax = plt.subplots(aug_number+1, 2, figsize=(8,aug_number*4))
show_img(ds[idx][0], ctx=ax[0,0])
ds[idx][1].show(ctx=ax[0,1])
# print(ds[idx][0])
for i in range(1, aug_number+1):
img, mask = aug.encodes(ds[idx])
show_img(img, ctx=ax[i,0])
mask.show(ctx=ax[i,1])
db = DataBlock(blocks=(TransformBlock([open_im]),
TransformBlock([partial(map_filename, str1='data', str2='label1'),
partial(open_im, cls=TensorMask)])),
splitter=RandomSplitter(valid_pct=0.2),
item_tfms=aug,
)
dl = db.dataloaders(items, bs=1)
idx=3
img, mask = dl.do_item(idx)
fig, (ax1, ax2) = plt.subplots(1,2)
tensor_img=img
ax1.imshow(tensor_img.permute(1, 2, 0)[..., :1])
ax2.imshow(mask)
#axs[1].imshow(mask)
plt.show()
print(np.shape(img))
print(np.shape(mask))
def acc_metric(input, target):
target = target.squeeze(1)
return (input.argmax(dim=1)==target).float().mean()
def loss_fn(pred, targ):
targ[targ==255] = 1
return torch.nn.functional.cross_entropy(pred, targ.squeeze(1).type(torch.long))
db = DataBlock(blocks=(TransformBlock([open_im]),
TransformBlock([partial(map_filename, str1='data', str2='label1'),
partial(open_im, cls=TensorMask)])),
splitter=RandomSplitter(valid_pct=0.2),
item_tfms=aug,
)
dl = db.dataloaders(items, bs=1)
learn = unet_learner(dls = dl, arch = resnet18, pretrained = True,normalize = False,n_in=2, n_out=2, loss_func=loss_fn, metrics=acc_metric)
learn.lr_find()
learn.fit_one_cycle(20,lr_max=6e-5, wd=0.8)
learn.fine_tune(8)
learn.export()
img,mask=dl.do_item(3)
cat, tensor, probs=learn.predict(items[1], masks[1])
cat, tensor, probs=learn.predict(img)
I tried to predict images in three different ways and also with learn.get_preds() and the dataloader, but it was not successful. The problem seems to be the encodes function for the masks, and images for the augmentation.
When i run: cat, tensor, probs=learn.predict(img)
The following error appears, but I don’t know how to fix this.
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
/tmp/ipykernel_14397/663310027.py in <module>
----> 1 cat, tensor, probs=learn.predict(img)
~/miniconda3/envs/fastai/lib/python3.9/site-packages/fastai/learner.py in predict(self, item, rm_type_tfms, with_input)
264 def predict(self, item, rm_type_tfms=None, with_input=False):
265 dl = self.dls.test_dl([item], rm_type_tfms=rm_type_tfms, num_workers=0)
--> 266 inp,preds,_,dec_preds = self.get_preds(dl=dl, with_input=True, with_decoded=True)
267 i = getattr(self.dls, 'n_inp', -1)
268 inp = (inp,) if i==1 else tuplify(inp)
~/miniconda3/envs/fastai/lib/python3.9/site-packages/fastai/learner.py in get_preds(self, ds_idx, dl, with_input, with_decoded, with_loss, act, inner, reorder, cbs, **kwargs)
251 if with_loss: ctx_mgrs.append(self.loss_not_reduced())
252 with ContextManagers(ctx_mgrs):
--> 253 self._do_epoch_validate(dl=dl)
254 if act is None: act = getattr(self.loss_func, 'activation', noop)
255 res = cb.all_tensors()
~/miniconda3/envs/fastai/lib/python3.9/site-packages/fastai/learner.py in _do_epoch_validate(self, ds_idx, dl)
201 if dl is None: dl = self.dls[ds_idx]
202 self.dl = dl
--> 203 with torch.no_grad(): self._with_events(self.all_batches, 'validate', CancelValidException)
204
205 def _do_epoch(self):
~/miniconda3/envs/fastai/lib/python3.9/site-packages/fastai/learner.py in _with_events(self, f, event_type, ex, final)
161
162 def _with_events(self, f, event_type, ex, final=noop):
--> 163 try: self(f'before_{event_type}'); f()
164 except ex: self(f'after_cancel_{event_type}')
165 self(f'after_{event_type}'); final()
~/miniconda3/envs/fastai/lib/python3.9/site-packages/fastai/learner.py in all_batches(self)
167 def all_batches(self):
168 self.n_iter = len(self.dl)
--> 169 for o in enumerate(self.dl): self.one_batch(*o)
170
171 def _do_one_batch(self):
~/miniconda3/envs/fastai/lib/python3.9/site-packages/fastai/data/load.py in __iter__(self)
107 self.before_iter()
108 self.__idxs=self.get_idxs() # called in context of main process (not workers/subprocesses)
--> 109 for b in _loaders[self.fake_l.num_workers==0](self.fake_l):
110 if self.device is not None: b = to_device(b, self.device)
111 yield self.after_batch(b)
~/miniconda3/envs/fastai/lib/python3.9/site-packages/torch/utils/data/dataloader.py in __next__(self)
519 if self._sampler_iter is None:
520 self._reset()
--> 521 data = self._next_data()
522 self._num_yielded += 1
523 if self._dataset_kind == _DatasetKind.Iterable and \
~/miniconda3/envs/fastai/lib/python3.9/site-packages/torch/utils/data/dataloader.py in _next_data(self)
559 def _next_data(self):
560 index = self._next_index() # may raise StopIteration
--> 561 data = self._dataset_fetcher.fetch(index) # may raise StopIteration
562 if self._pin_memory:
563 data = _utils.pin_memory.pin_memory(data)
~/miniconda3/envs/fastai/lib/python3.9/site-packages/torch/utils/data/_utils/fetch.py in fetch(self, possibly_batched_index)
32 raise StopIteration
33 else:
---> 34 data = next(self.dataset_iter)
35 return self.collate_fn(data)
36
~/miniconda3/envs/fastai/lib/python3.9/site-packages/fastai/data/load.py in create_batches(self, samps)
116 if self.dataset is not None: self.it = iter(self.dataset)
117 res = filter(lambda o:o is not None, map(self.do_item, samps))
--> 118 yield from map(self.do_batch, self.chunkify(res))
119
120 def new(self, dataset=None, cls=None, **kwargs):
~/miniconda3/envs/fastai/lib/python3.9/site-packages/fastcore/basics.py in chunked(it, chunk_sz, drop_last, n_chunks)
214 if not isinstance(it, Iterator): it = iter(it)
215 while True:
--> 216 res = list(itertools.islice(it, chunk_sz))
217 if res and (len(res)==chunk_sz or not drop_last): yield res
218 if len(res)<chunk_sz: return
~/miniconda3/envs/fastai/lib/python3.9/site-packages/fastai/data/load.py in do_item(self, s)
131 def prebatched(self): return self.bs is None
132 def do_item(self, s):
--> 133 try: return self.after_item(self.create_item(s))
134 except SkipItemException: return None
135 def chunkify(self, b): return b if self.prebatched else chunked(b, self.bs, self.drop_last)
~/miniconda3/envs/fastai/lib/python3.9/site-packages/fastcore/transform.py in __call__(self, o)
198 self.fs = self.fs.sorted(key='order')
199
--> 200 def __call__(self, o): return compose_tfms(o, tfms=self.fs, split_idx=self.split_idx)
201 def __repr__(self): return f"Pipeline: {' -> '.join([f.name for f in self.fs if f.name != 'noop'])}"
202 def __getitem__(self,i): return self.fs[i]
~/miniconda3/envs/fastai/lib/python3.9/site-packages/fastcore/transform.py in compose_tfms(x, tfms, is_enc, reverse, **kwargs)
148 for f in tfms:
149 if not is_enc: f = f.decode
--> 150 x = f(x, **kwargs)
151 return x
152
~/miniconda3/envs/fastai/lib/python3.9/site-packages/fastcore/transform.py in __call__(self, x, **kwargs)
111 "A transform that always take tuples as items"
112 _retain = True
--> 113 def __call__(self, x, **kwargs): return self._call1(x, '__call__', **kwargs)
114 def decode(self, x, **kwargs): return self._call1(x, 'decode', **kwargs)
115 def _call1(self, x, name, **kwargs):
~/miniconda3/envs/fastai/lib/python3.9/site-packages/fastcore/transform.py in _call1(self, x, name, **kwargs)
115 def _call1(self, x, name, **kwargs):
116 if not _is_tuple(x): return getattr(super(), name)(x, **kwargs)
--> 117 y = getattr(super(), name)(list(x), **kwargs)
118 if not self._retain: return y
119 if is_listy(y) and not isinstance(y, tuple): y = tuple(y)
~/miniconda3/envs/fastai/lib/python3.9/site-packages/fastcore/transform.py in __call__(self, x, **kwargs)
71 #property
72 def name(self): return getattr(self, '_name', _get_name(self))
---> 73 def __call__(self, x, **kwargs): return self._call('encodes', x, **kwargs)
74 def decode (self, x, **kwargs): return self._call('decodes', x, **kwargs)
75 def __repr__(self): return f'{self.name}:\nencodes: {self.encodes}decodes: {self.decodes}'
~/miniconda3/envs/fastai/lib/python3.9/site-packages/fastcore/transform.py in _call(self, fn, x, split_idx, **kwargs)
81 def _call(self, fn, x, split_idx=None, **kwargs):
82 if split_idx!=self.split_idx and self.split_idx is not None: return x
---> 83 return self._do_call(getattr(self, fn), x, **kwargs)
84
85 def _do_call(self, f, x, **kwargs):
~/miniconda3/envs/fastai/lib/python3.9/site-packages/fastcore/transform.py in _do_call(self, f, x, **kwargs)
87 if f is None: return x
88 ret = f.returns(x) if hasattr(f,'returns') else None
---> 89 return retain_type(f(x, **kwargs), x, ret)
90 res = tuple(self._do_call(f, x_, **kwargs) for x_ in x)
91 return retain_type(res, x)
~/miniconda3/envs/fastai/lib/python3.9/site-packages/fastcore/dispatch.py in __call__(self, *args, **kwargs)
116 elif self.inst is not None: f = MethodType(f, self.inst)
117 elif self.owner is not None: f = MethodType(f, self.owner)
--> 118 return f(*args, **kwargs)
119
120 def __get__(self, inst, owner):
/tmp/ipykernel_14397/3758110305.py in encodes(self, x)
7
8 def encodes(self, x):
----> 9 img,mask = x
10 img = img/img.max()
11 aug = self.aug(image=np.array(img.permute(1,2,0)), mask=np.array(mask))
ValueError: not enough values to unpack (expected 2, got 1)
Also cat, tensor, probs=learn.predict(items1, masks1) gives the error:
TypeError: slice indices must be integers or None or have an __index__ method
But I don't know why.