med-air/Endo-FM

HOW TO SEE OUTPUT AFTER DIRECT DOWNSTREAM TESTING

prabin333 opened this issue · 8 comments

~/Endo-FM/TransUNet$ python train.py --test
Namespace(base_lr=0.0001, batch_size=1, dataset='Synapse', deterministic=1, exp='TU_Synapse224', img_size=224, is_pretrain=True, list_dir='./lists/lists_Synapse', max_epochs=150, max_iterations=30000, n_gpu=1, n_skip=3, num_classes=2, pretrained_model_weights='cvc.pth', root_path='../data/pretrain/CVC-ClinicVideoDB/', seed=9041, test=True, vit_name='R50-ViT-B_16', vit_patches_size=16)
The length of train set is: 20
The length of test set is: 9
Test Dice: 76.6, HD95: 7.0

Hi, thanks for your interest!
May you elaborate more on your question, about what ouput?

I used the CVC segmentation model for direct downstream testing. I ran the code you gave me without any errors. I pasted the code and output in the previous comment, but i don't know the output is stored or not . There is no clear information about the output directory. Can you please explain how to locate the output directory?

~/Endo-FM/TransUNet$ python train.py --test
Namespace(base_lr=0.0001, batch_size=1, dataset='Synapse', deterministic=1, exp='TU_Synapse224', img_size=224, is_pretrain=True, list_dir='./lists/lists_Synapse', max_epochs=150, max_iterations=30000, n_gpu=1, n_skip=3, num_classes=2, pretrained_model_weights='cvc.pth', root_path='../data/pretrain/CVC-ClinicVideoDB/', seed=9041, test=True, vit_name='R50-ViT-B_16', vit_patches_size=16)
The length of train set is: 20
The length of test set is: 9
Test Dice: 76.6, HD95: 7.0

Hi, we do not save the predictions from the model in this repo, but just output the metrics.
You can use argmax to this output and then save the prediction:

output = model(img)

Could you please provide a step-by-step process with code examples for my query? I would like to test it with a video file. My aim is to obtain the same output as the one demonstrated in your GitHub account. Please let me know if there are any additional steps I need to follow beyond the ones mentioned in your GitHub documentation.

Hi, here is a sample for getting predictions:

def inference(model, val_loader, device):
    model.eval()
    savedir = 'saved/endofm'
    os.makedirs(savedir, exist_ok=True)

    with torch.no_grad():
        for batch in tqdm(val_loader):
            img, label = batch['image'].to(device).squeeze(0), batch['label'].to(device).squeeze(0)
            video_info = batch['video_info']
            start_idx = int(video_info[1][0])
            end_idx = start_idx + int(video_info[2][0])

            output = model(img)
            prediction = torch.argmax(output, dim=1).cpu().numpy() * 255

            for i in range(start_idx, end_idx):
                image_name = f'{i}.tif'
                cv2.imwrite(os.path.join(savedir, image_name), prediction[i - start_idx])

if args.inference:
    model.eval()
    model.load_state_dict(torch.load(args.pretrained_model_weights, map_location='cpu'))
    inference(model, testloader, 'cuda')
    exit(0)

You can just integrate this after this test args:

if args.test:
model.eval()
model.load_state_dict(torch.load(args.pretrained_model_weights, map_location='cpu'))
test_dice, test_hd95 = eval(model, testloader, 'cuda', classes=2)
print('Test Dice: %.1f, HD95: %.1f' % (test_dice * 100., test_hd95))
exit(0)

HI, thanks for all you great work!!!
I used CVC database for directing downstream test, and I also want to check out outputs in a folder. I modified the code in the "def eval()"section in trainer.py:

def eval(model, val_loader, device, classes, output_folder):
    all_dice = []
    all_hd95 = []
    model.eval()
    with torch.no_grad():
        for i, batch in enumerate(val_loader):
            img, label = batch['image'].to(device).squeeze(0), batch['label'].to(device).squeeze(0)
            output = model(img)
            save_output_images(output, output_folder, i)
            dice, hd95 = eval_dice(output, label, classes=classes)
            all_dice.append(dice.item())
            all_hd95.append(hd95.item())
    return np.mean(np.array(all_dice)), np.mean(np.array(all_hd95))

def save_output_images(output, output_folder, batch_idx):
    os.makedirs(output_folder, exist_ok=True)
    for i, img_tensor in enumerate(output):
        img_pil = TF.to_pil_image(img_tensor.cpu())
        img_pil.save(os.path.join(output_folder, f'output_batch_{batch_idx}_img_{i}.png'))
        
if args.test:
    model.eval()
    model.load_state_dict(torch.load(args.pretrained_model_weights, map_location='cpu'))
    test_dice, test_hd95 = eval(model, testloader, 'cuda', classes=2, output_folder='/root/test0806/Endo-FM/TransUNet/saved/endofm/')
    print('Test Dice: %.1f, HD95: %.1f' % (test_dice * 100., test_hd95))
    exit(0)

then I got strange outputs like these:
output_batch_0_img_0
image
then I changed some data pics and modified the suffix, still unresolved.

Hi, @Jiangmiemiea
Thanks for your interest! The original output from the model is not the predicted segmentation result, you should first do this to obtain the predicted mask (* 255 is used to recover the image to [0, 255]):

prediction = torch.argmax(output, dim=1).cpu().numpy() * 255

您好,感谢您的关注!模型的原始输出不是预测的分割结果,您应该首先这样做以获得预测的掩码(用于将图像恢复到 [0, 255]):* 255

prediction = torch.argmax(output, dim=1).cpu().numpy() * 255

Thanks! Problem solved. :)