Как получить правильные выходные прогнозы от unet_learner (fastai)? - PullRequest
Новая
       121

Как получить правильные выходные прогнозы от unet_learner (fastai)?

1 голос
/ 15 января 2020

Пожалуйста, я работаю над проектом сегментации изображений и использовал библиотеку 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))

enter image description here

Пожалуйста, есть какие-нибудь идеи о том, как решить эту проблему? Спасибо

Ответы [ 2 ]

0 голосов
/ 02 мая 2020

Чтобы просмотреть прогноз unet_learner с наложением поверх исходного изображения, вы можете сделать: 

img = open_image("your_test_image.png")
prediction = learn.predict(img)
img.show(y=prediction[0])

Вот пример в документах fastai: https://docs.fast.ai/tutorial.inference.html#A -segmentation-example .

0 голосов
/ 20 января 2020

Я думаю, что прогноз в порядке. Вы ожидаете что-то подобное?

prediction result

Этот результат основан на опубликованном вами прогнозном изображении.

Чтобы проверить, как идут дела, попробуйте следующее: 

 interp = SegmentationInterpretation.from_learner(learn)
 mean_cm, single_img_cm = interp._generate_confusion()
 df = interp._plot_intersect_cm(mean_cm, "Mean of Ratio of Intersection given 
 True Label")
 i = 0 #Some image index
 df = interp._plot_intersect_cm(single_img_cm[i], f"Ratio of Intersection given True Label, Image:{i}")
 interp.show_xyz(i)

Основано на документах fast.ai

Что касается результата прогнозирования, то это изображение основано на значениях ваших классов. Если вы возьмете значения (r, g, b) из этого изображения, у вас будет (r, g, b) == 0 для фона и (r, g, b) == 1 для краев. Если у вас есть больше классов, следующий будет как (r, g, b) == 2 и так далее.

Так что вы можете просто раскрасить свой результат прогноза. Я сделал это с помощью OpenCV, что-то вроде этого: 

  frame = cv2.imread("yourPredictionHere.png",1)
  frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB) 
  for x in range(384): #width based on the size of your image.
      for y in range(384): #height based on the size of your image.
          b, g, r = frame[x, y]
          if (b, g, r) == (0,0,0): #background
              frame[x, y] = (0,0,0)
          elif (b, g, r) == (1,1,1): #edges
              frame[x, y] = (85,85,255)

  cv2.imwrite("result.png",frame)

С уважением!

...