| 303 | |
| 304 | |
| 305 | def vis_image(imgs, pred_masks, gt_masks, save_path, reverse = False, points = None): |
| 306 | |
| 307 | b,c,h,w = pred_masks.size() |
| 308 | dev = pred_masks.get_device() |
| 309 | row_num = min(b, 4) |
| 310 | |
| 311 | if torch.max(pred_masks) > 1 or torch.min(pred_masks) < 0: |
| 312 | pred_masks = torch.sigmoid(pred_masks) |
| 313 | |
| 314 | if reverse == True: |
| 315 | pred_masks = 1 - pred_masks |
| 316 | gt_masks = 1 - gt_masks |
| 317 | if c == 2: # for REFUGE multi mask output |
| 318 | pred_disc, pred_cup = pred_masks[:,0,:,:].unsqueeze(1).expand(b,3,h,w), pred_masks[:,1,:,:].unsqueeze(1).expand(b,3,h,w) |
| 319 | gt_disc, gt_cup = gt_masks[:,0,:,:].unsqueeze(1).expand(b,3,h,w), gt_masks[:,1,:,:].unsqueeze(1).expand(b,3,h,w) |
| 320 | tup = (imgs[:row_num,:,:,:],pred_disc[:row_num,:,:,:], pred_cup[:row_num,:,:,:], gt_disc[:row_num,:,:,:], gt_cup[:row_num,:,:,:]) |
| 321 | compose = torch.cat(tup, 0) |
| 322 | vutils.save_image(compose, fp = save_path, nrow = row_num, padding = 10) |
| 323 | elif c > 2: # for multi-class segmentation > 2 classes |
| 324 | preds = [] |
| 325 | gts = [] |
| 326 | for i in range(0, c): |
| 327 | pred = pred_masks[:,i,:,:].unsqueeze(1).expand(b,3,h,w) |
| 328 | preds.append(pred) |
| 329 | gt = gt_masks[:,i,:,:].unsqueeze(1).expand(b,3,h,w) |
| 330 | gts.append(gt) |
| 331 | tup = [imgs[:row_num,:,:,:]] + preds + gts |
| 332 | compose = torch.cat(tup,0) |
| 333 | vutils.save_image(compose, fp = save_path, nrow = row_num, padding = 10) |
| 334 | else: |
| 335 | imgs = torchvision.transforms.Resize((h,w))(imgs) |
| 336 | if imgs.size(1) == 1: |
| 337 | imgs = imgs[:,0,:,:].unsqueeze(1).expand(b,3,h,w) |
| 338 | pred_masks = pred_masks[:,0,:,:].unsqueeze(1).expand(b,3,h,w) |
| 339 | gt_masks = gt_masks[:,0,:,:].unsqueeze(1).expand(b,3,h,w) |
| 340 | if points != None: |
| 341 | for i in range(b): |
| 342 | |
| 343 | p = np.round(points.cpu()/args.image_size * args.out_size).to(dtype = torch.int) |
| 344 | |
| 345 | gt_masks[i,0,p[i,0]-2:p[i,0]+2,p[i,1]-2:p[i,1]+2] = 0.5 |
| 346 | gt_masks[i,1,p[i,0]-2:p[i,0]+2,p[i,1]-2:p[i,1]+2] = 0.1 |
| 347 | gt_masks[i,2,p[i,0]-2:p[i,0]+2,p[i,1]-2:p[i,1]+2] = 0.4 |
| 348 | # gt_masks[i,0,p[i,0]-5:p[i,0]+5,p[i,1]-5:p[i,1]+5] = 0.5 |
| 349 | # gt_masks[i,1,p[i,0]-5:p[i,0]+5,p[i,1]-5:p[i,1]+5] = 0.1 |
| 350 | # gt_masks[i,2,p[i,0]-5:p[i,0]+5,p[i,1]-5:p[i,1]+5] = 0.4 |
| 351 | tup = (imgs[:row_num,:,:,:],pred_masks[:row_num,:,:,:], gt_masks[:row_num,:,:,:]) |
| 352 | # compose = torch.cat((imgs[:row_num,:,:,:],pred_disc[:row_num,:,:,:], pred_cup[:row_num,:,:,:], gt_disc[:row_num,:,:,:], gt_cup[:row_num,:,:,:]),0) |
| 353 | compose = torch.cat(tup,0) |
| 354 | vutils.save_image(compose, fp = save_path, nrow = row_num, padding = 10) |
| 355 | |
| 356 | return |
| 357 | |
| 358 | def eval_seg(pred,true_mask_p,threshold): |
| 359 | ''' |