#machine-learning #computer-vision #fast-ai #unity3d-unet
Вопрос:
У меня есть приведенный ниже код для обучения модели и сохранения ее как model_7.pth с помощью fastai.
from fastai import * from wwf.utils import state_versions state_versions(['fastai', 'fastcore', 'wwf']) from fastai.vision.all import * path = Path('data') def n_codes(fnames, is_partial=True): "Gather the codes from a list of `fnames`" vals = set() if is_partial: random.shuffle(fnames) fnames = fnames[:10] for fname in fnames: msk = np.array(PILMask.create(fname)) for val in np.unique(msk): if val not in vals: vals.add(val) vals = list(vals) p2c = dict() for i,val in enumerate(vals): p2c[i] = vals[i] return p2c p2c = n_codes(lbl_names) def get_msk(fn, pix2class): "Grab a mask from a `filename` and adjust the pixels based on `pix2class`" #fn = path/'GT_png'/f'{fn.stem}_mask.png' fn = path/f'trainval/train/tile/label/{fn.name}' msk = np.array(PILMask.create(fn)) mx = np.max(msk) for i, val in enumerate(p2c): msk[msk==p2c[i]] = val return PILMask.create(msk) codes = ['Background', 'car'] get_y = lambda o: get_msk(o, p2c) binary = DataBlock(blocks=(ImageBlock, MaskBlock(codes)), get_items=get_image_files, splitter=RandomSplitter(), get_y=get_y, item_tfms=Resize(224), batch_tfms=[Normalize.from_stats(*imagenet_stats)]) dls = binary.dataloaders(path/'trainval/train/tile/label', bs=4, num_workers=0)#num_workers=0 to make it work in windows dls.show_batch(cmap='Blues', vmin=0, vmax=1) learn = unet_learner(dls, resnet34) learn.load(r'model_7') learn.fit(n_epoch=8, cbs=SaveModelCallback(every_epoch=True)) Теперь у меня есть папка с тестовым изображением под названием TEST_FOLDER.
Я попробовал приведенный ниже код.
TEST_FOLDER = r"testtileimage" files = get_image_files(TEST_FOLDER)[:1] test_dl = learn.dls.test_dl(files) preds = learn.get_preds(dl=test_dl) но выходные данные имеют форму
preds[0].shape Out[43]: torch.Size([1, 2, 224, 224]) Я использую этот код, чтобы увидеть прогноз ожидаемого результата
p = preds[0][0] plt.plot(p[0]) plt.savefig('myfig_p0') но я получаю изображение с 4 подзаголовками и не уверен, какая из них является предсказанной маской. Кроме того, нижнее левое изображение не соответствует входному изображению. 
Вопрос
- Как я могу предсказать маску сегмента для каждого тестового изображения?
- Как получить пиксельную вероятность для каждого тестового изображения?