OVSegmentor/main_demo.py

# -------------------------------------------------------------------------
# Written by Jilan Xu
# -------------------------------------------------------------------------

import argparse
import os
import os.path as osp
import subprocess

import mmcv
import torch
import torch.backends.cudnn as cudnn
import torch.distributed as dist
from datasets import build_text_transform
from main_pretrain import validate_seg
from mmcv.image import tensor2imgs
from mmcv.parallel import MMDistributedDataParallel
from mmcv.runner import set_random_seed
from models import build_model
from omegaconf import OmegaConf, read_write

from segmentation.evaluation import build_seg_dataloader, build_seg_dataset, build_custom_seg_dataset, build_seg_inference, build_demo_inference
from utils import get_config, get_logger, load_checkpoint
from transformers import AutoTokenizer, RobertaTokenizer
from ipdb import set_trace
from main_pretrain import init_distributed_mode

try:
    # noinspection PyUnresolvedReferences
    from apex import amp
except ImportError:
    amp = None
    
tokenizer_dict = {
    'Bert': AutoTokenizer.from_pretrained('distilbert-base-uncased', TOKENIZERS_PARALLELISM=False),
    # 'Roberta': RobertaTokenizer.from_pretrained('/mnt/petrelfs/xujilan/roberta-base/'),
    'Roberta': RobertaTokenizer.from_pretrained('roberta-base'),
    'TextTransformer': None,
}

def parse_args():
    parser = argparse.ArgumentParser('OVSegmentor segmentation demo')
    parser.add_argument(
        '--cfg',
        type=str,
        default='./configs/test_voc12.yml',
        help='path to config file',
    )
    parser.add_argument(
        '--resume', 
        type=str,
        required=True,
        help='resume from checkpoint',
    )
    parser.add_argument(
        '--image_folder',
        type=str,
        required=True,
        help='path to the input image folder',
    )
    parser.add_argument(
        '--vocab',
        help='could be a list of candidate vocabularies, use given classes from [voc, coco, ade], or give a custom list of classes',
        default='voc',
        nargs='+',
    )
    
    parser.add_argument(
        '--output_folder', 
        type=str, 
        help='root of output folder',
    )
    parser.add_argument(
        '--vis',
        help='Specify the visualization mode, '
        'could be a list, support input, pred, input_seg, input_pred_seg_label, all_groups, first_group, last_group, mask',
        default='input_pred_seg_label',
        nargs='+',
    )
    parser.add_argument(
        '--opts',
        help="Modify config options by adding 'KEY VALUE' pairs. ",
        default=None,
        nargs='+',
    )


    # distributed training
    parser.add_argument('--local_rank', type=int, required=False, default=0, help='local rank for DistributedDataParallel')
    args = parser.parse_args()
    return args

def generate_imagelist_with_sanity_check(root):
    image_list = []
    for each_file in os.listdir(root):
        ### assume we process all .jpg files, and convert png to jpg files
        if each_file.endswith('.jpg'):
            pass
        elif each_file.endswith('.png'):
            img = mmcv.imread(osp.join(root, each_file))
            mmcv.imwrite(img, osp.join(root, each_file.replace('.png','.jpg')))
        else:
            continue
        
        filename = each_file.split('.')[0]
        image_list.append(filename)
            
    if len(image_list) == 0:
        raise ValueError(f'No image found in {args.image_folder}')
    
    with open(os.path.join(root, 'image_list.txt'), 'w') as f:
        for item in image_list:
            f.write("%s\n" % item)
            
    return image_list
    
def inference(cfg):
    logger = get_logger()
    
    ### check and generate image list ###
    generate_imagelist_with_sanity_check(cfg.image_folder)
    os.makedirs(cfg.output_folder, exist_ok=True)
    
    data_loader = build_seg_dataloader(build_custom_seg_dataset(cfg.evaluate.seg, cfg))
    dataset = data_loader.dataset
    print('whether activating visualization: ', cfg.vis)
    logger.info(f'Evaluating dataset: {dataset}')
    
    logger.info(f'Creating model:{cfg.model.type}/{cfg.model_name}')
    model = build_model(cfg.model)
    model.cuda()
    logger.info(str(model))

    if cfg.train.amp_opt_level != 'O0':
        model = amp.initialize(model, None, opt_level=cfg.train.amp_opt_level)

    n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
    logger.info(f'number of params: {n_parameters}')

    load_checkpoint(cfg, model, None, None)
    
    global tokenizer
    tokenizer = tokenizer_dict[cfg.model.text_encoder.type]

    if cfg.vis:
        vis_seg(cfg, data_loader, model, cfg.vis)


@torch.no_grad()
def vis_seg(config, data_loader, model, vis_modes):
    dist.barrier()
    model.eval()
    
    if hasattr(model, 'module'):
        model_without_ddp = model.module
    else:
        model_without_ddp = model

    text_transform = build_text_transform(False, config.data.text_aug, with_dc=False)
    seg_model = build_demo_inference(model_without_ddp, text_transform, config, tokenizer)

    mmddp_model = MMDistributedDataParallel(
        seg_model, device_ids=[torch.cuda.current_device()], broadcast_buffers=False)
    mmddp_model.eval()
    model = mmddp_model.module
    device = next(model.parameters()).device
    dataset = data_loader.dataset
    
    if dist.get_rank() == 0:
        prog_bar = mmcv.ProgressBar(len(dataset))
    loader_indices = data_loader.batch_sampler
    for batch_indices, data in zip(loader_indices, data_loader):
        with torch.no_grad():
            result = mmddp_model(return_loss=False, **data)
        # set_trace()

        img_tensor = data['img'][0]
        img_metas = data['img_metas'][0].data[0]
        imgs = tensor2imgs(img_tensor, **img_metas[0]['img_norm_cfg'])
        assert len(imgs) == len(img_metas)

        for batch_idx, img, img_meta in zip(batch_indices, imgs, img_metas):
            h, w, _ = img_meta['img_shape']
            img_show = img[:h, :w, :]

            ori_h, ori_w = img_meta['ori_shape'][:-1]
            img_show = mmcv.imresize(img_show, (ori_w, ori_h))

            for vis_mode in vis_modes:
                out_file = osp.join(config.output_folder, vis_mode, f'{batch_idx:04d}.jpg')
                # os.makedirs(osp.join(config.output_folder, 'vis_imgs', vis_mode), exist_ok=True)
                print(osp.join(config.output_folder, vis_mode))
                model.show_result(img_show, img_tensor.to(device), result, out_file, vis_mode)
            if dist.get_rank() == 0:
                batch_size = len(result) * dist.get_world_size()
                for _ in range(batch_size):
                    prog_bar.update()    

def main():
    args = parse_args()
    cfg = get_config(args)

    if cfg.train.amp_opt_level != 'O0':
        assert amp is not None, 'amp not installed!'

    with read_write(cfg):
        cfg.evaluate.eval_only = True
    
    init_distributed_mode(args)
    rank, world_size = args.rank, args.world_size

    set_random_seed(cfg.seed, use_rank_shift=True)
    cudnn.benchmark = True

    os.makedirs(cfg.output, exist_ok=True)
    logger = get_logger(cfg)

    if dist.get_rank() == 0:
        path = os.path.join(cfg.output, 'config.json')
        OmegaConf.save(cfg, path)
        logger.info(f'Full config saved to {path}')

    # print config
    logger.info(OmegaConf.to_yaml(cfg))

    inference(cfg)
    dist.barrier()


if __name__ == '__main__':
    main()