threshold: a int or a tuple of int masks: [b,2,h,w] pred: [b,2,h,w]
(pred,true_mask_p,threshold)
| 356 | return |
| 357 | |
| 358 | def eval_seg(pred,true_mask_p,threshold): |
| 359 | ''' |
| 360 | threshold: a int or a tuple of int |
| 361 | masks: [b,2,h,w] |
| 362 | pred: [b,2,h,w] |
| 363 | ''' |
| 364 | b, c, h, w = pred.size() |
| 365 | if c == 2: |
| 366 | iou_d, iou_c, disc_dice, cup_dice = 0,0,0,0 |
| 367 | for th in threshold: |
| 368 | |
| 369 | gt_vmask_p = (true_mask_p > th).float() |
| 370 | vpred = (pred > th).float() |
| 371 | vpred_cpu = vpred.cpu() |
| 372 | disc_pred = vpred_cpu[:,0,:,:].numpy().astype('int32') |
| 373 | cup_pred = vpred_cpu[:,1,:,:].numpy().astype('int32') |
| 374 | |
| 375 | disc_mask = gt_vmask_p [:,0,:,:].squeeze(1).cpu().numpy().astype('int32') |
| 376 | cup_mask = gt_vmask_p [:, 1, :, :].squeeze(1).cpu().numpy().astype('int32') |
| 377 | |
| 378 | '''iou for numpy''' |
| 379 | iou_d += iou(disc_pred,disc_mask) |
| 380 | iou_c += iou(cup_pred,cup_mask) |
| 381 | |
| 382 | '''dice for torch''' |
| 383 | disc_dice += dice_coeff(vpred[:,0,:,:], gt_vmask_p[:,0,:,:]).item() |
| 384 | cup_dice += dice_coeff(vpred[:,1,:,:], gt_vmask_p[:,1,:,:]).item() |
| 385 | |
| 386 | return iou_d / len(threshold), iou_c / len(threshold), disc_dice / len(threshold), cup_dice / len(threshold) |
| 387 | elif c > 2: # for multi-class segmentation > 2 classes |
| 388 | ious = [0] * c |
| 389 | dices = [0] * c |
| 390 | for th in threshold: |
| 391 | gt_vmask_p = (true_mask_p > th).float() |
| 392 | vpred = (pred > th).float() |
| 393 | vpred_cpu = vpred.cpu() |
| 394 | for i in range(0, c): |
| 395 | pred = vpred_cpu[:,i,:,:].numpy().astype('int32') |
| 396 | mask = gt_vmask_p[:,i,:,:].squeeze(1).cpu().numpy().astype('int32') |
| 397 | |
| 398 | '''iou for numpy''' |
| 399 | ious[i] += iou(pred,mask) |
| 400 | |
| 401 | '''dice for torch''' |
| 402 | dices[i] += dice_coeff(vpred[:,i,:,:], gt_vmask_p[:,i,:,:]).item() |
| 403 | |
| 404 | return tuple(np.array(ious + dices) / len(threshold)) # tuple has a total number of c * 2 |
| 405 | else: |
| 406 | eiou, edice = 0,0 |
| 407 | for th in threshold: |
| 408 | |
| 409 | gt_vmask_p = (true_mask_p > th).float() |
| 410 | vpred = (pred > th).float() |
| 411 | vpred_cpu = vpred.cpu() |
| 412 | disc_pred = vpred_cpu[:,0,:,:].numpy().astype('int32') |
| 413 | |
| 414 | disc_mask = gt_vmask_p [:,0,:,:].squeeze(1).cpu().numpy().astype('int32') |
| 415 |
no test coverage detected