Пожалуйста, я работаю над проектом сегментации изображений и использовал библиотеку fastai (в частности, unet_learner). Я обучил свою модель, и все в порядке, вот мой код (на этапе обучения):
#codes = np.loadtxt('codes.txt', dtype=str)
codes = np.array(['bg', 'edge'], dtype='<U4')# bg= background
get_y_fn = lambda x: path_lbl/f'{x.stem}{x.suffix}'
# fastai codes
data = (SegmentationItemList.from_folder(path_img)
.split_by_rand_pct()
.label_from_func(get_y_fn, classes=codes)
#.add_test_folder()
#.transform(get_transforms(), tfm_y=True, size=384)
.databunch(bs=2,path=dataset) # bs = mimi-patch size
.normalize(imagenet_stats))
learn = unet_learner(data, models.resnet34, wd=1e-2)
learn.lr_find() # find learning rate
learn.recorder.plot() # plot learning rate graph
lr = 1e-02 # pick a lr
learn.fit_one_cycle(3, slice(lr), pct_start=0.3) # train model ---- epochs=3
learn.unfreeze() # unfreeze all layers
# find and plot lr again
learn.lr_find()
learn.recorder.plot()
learn.fit_one_cycle(10, slice(lr/400, lr/4), pct_start=0.3)
learn.save('model-stage-1') # save model
learn.load('model-stage-1');
learn.export()
Моя проблема в том, что когда я пытаюсь сделать прогноз, используя обученную модель, на выходе всегда черное изображение. Ниже приведен код на этапе прогнозирования:
img = open_image('/content/generated_samples_masks/545.png')
prediction = learn.predict(img)
prediction[0].show(figsize=(8,8))
Пожалуйста, есть идеи, как решить эту проблему? Спасибо