M3d_eval_cap.py

import argparse
import json
import os

import mmengine
import numpy as np
from PIL import Image
import torch.nn as nn
import torch
import torch.distributed
import torch.utils.data
import tqdm
from transformers import AutoModel, AutoTokenizer
from projects.llava_sam2.evaluation.dataset import M3DSegDataset_eval, M3DVQADataset_closed_eval, M3DCAPDataset_eval
from projects.llava_sam2.evaluation.utils import _init_dist_pytorch, _init_dist_slurm, get_dist_info, get_rank, \
    collect_results_cpu

import concurrent.futures
from pycocotools import mask as cocomask


def async_func(executor, func, **kwargs):
    future = executor.submit(func, **kwargs)
    return future


def mask_to_rle(mask):
    rle = []
    for m in mask:
        rle.append(cocomask.encode(np.asfortranarray(m.astype(np.uint8))))
        rle[-1]['counts'] = rle[-1]['counts'].decode()
    return rle


def mask_save(item, mask_prediction, work_dir):
    vid_id = item['video_id']
    exp_id = item['exp_id']
    save_path = os.path.join(work_dir, 'Annotations', vid_id, exp_id)
    mmengine.mkdir_or_exist(save_path)
    for id_m, mask in enumerate(mask_prediction):
        mask = Image.fromarray(mask.astype(np.float32) * 255).convert('L')
        file_name = item['frames'][id_m]
        save_file = os.path.join(save_path, file_name + ".png")
        mask.save(save_file)


def dice_score(pred, target, epsilon=1e-6):
    """
    Compute Dice score between two binary 3D images.

    Args:
        pred (np.ndarray): Prediction array of shape (D, H, W)
        target (np.ndarray): Ground truth array of shape (D, H, W)
        epsilon (float): Small value to avoid division by zero

    Returns:
        float: Dice coefficient
    """
    pred = pred.astype(bool)
    target = target.astype(bool)

    intersection = np.logical_and(pred, target).sum()
    dice = (2. * intersection) / (pred.sum() + target.sum() + epsilon)
    return dice

DATASETS_INFO = {
    'DAVIS': {
        'data_root': 'data/video_datas/davis17/',
        'image_folder': 'data/video_datas/davis17/valid/JPEGImages/',
        'expression_file': 'data/video_datas/davis17/meta_expressions/valid/meta_expressions.json',
        'mask_file': 'data/video_datas/davis17/valid/mask_dict.pkl',
    },
    'MEVIS': {
        'data_root': 'data/video_datas/mevis/valid/',
        'image_folder': 'data/video_datas/mevis/valid/JPEGImages',
        'expression_file': 'data/video_datas/mevis/valid/meta_expressions.json',
        'mask_file': None,
    },
    'MEVIS_U': {
        'data_root': 'data/video_datas/mevis/valid_u/',
        'image_folder': 'data/video_datas/mevis/valid_u/JPEGImages',
        'expression_file': 'data/video_datas/mevis/valid_u/meta_expressions.json',
        'mask_file': 'data/video_datas/mevis/valid_u/mask_dict.json',
    },
    'REFYTVOS': {
        'data_root': 'data/video_datas/rvos/',
        'image_folder': 'data/video_datas/rvos/valid/JPEGImages/',
        'expression_file': 'data/video_datas/rvos/meta_expressions/valid/meta_expressions.json',
        'mask_file': None,
    },
    'REVOS': {
        'data_root': 'data/video_datas/revos/',
        'image_folder': 'data/video_datas/revos/',
        'expression_file': 'data/video_datas/revos/meta_expressions_valid_.json',
        'mask_file': None,},

    'ACT-1K': {
        'data_root': 'data/M3D_Seg/',
        'image_folder': 'data/M3D_Seg/M3D-Seg',
        'expression_file': '0003',
        'mask_file': None,
    },

    'TS': {
        'data_root': 'data/M3D_Seg/',
        'image_folder': 'data/M3D_Seg/M3D-Seg',
        'expression_file': '0011',
        'mask_file': None,
    },

    'CTOrg': {
        'data_root': 'data/M3D_Seg/',
        'image_folder': 'data/M3D_Seg/M3D-Seg',
        'expression_file': '0012',
        'mask_file': None,
    },
}


def parse_args():
    parser = argparse.ArgumentParser(description='RefVOS')
    parser.add_argument('model_path', help='hf model path.')
    parser.add_argument(
        '--vqa_file',
        default='/blue/kgong/yxing2/Sa2VA-main/data/M3D_VQA/M3D_VQA_test5k.csv',
        help='Specify a dataset')
    parser.add_argument(
        '--launcher',
        choices=['none', 'pytorch', 'slurm', 'mpi'],
        default='none',
        help='job launcher')
    parser.add_argument('--local_rank', '--local-rank', type=int, default=0)
    parser.add_argument('--submit', action='store_true')
    parser.add_argument('--work_dir', type=str, default=None)
    parser.add_argument('--deepspeed', type=str, default=None)  # dummy
    args = parser.parse_args()
    if 'LOCAL_RANK' not in os.environ:
        os.environ['LOCAL_RANK'] = str(args.local_rank)
    return args


if __name__ == '__main__':
    args = parse_args()

    work_dir = args.work_dir
    if work_dir is None:
        work_dir = 'work_dirs/foobar'

    if args.launcher == 'none':
        rank = 0
        world_size = 1
    elif args.launcher == 'pytorch':
        _init_dist_pytorch('nccl')
        rank, world_size = get_dist_info()
    elif args.launcher == 'slurm':
        _init_dist_slurm('nccl')
        rank, world_size = get_dist_info()

    model = AutoModel.from_pretrained(
        args.model_path,
        torch_dtype=torch.bfloat16,
        use_flash_attn=True,
        trust_remote_code=True,
    ).eval().cuda()

    tokenizer = AutoTokenizer.from_pretrained(
        args.model_path,
        trust_remote_code=True,
    )

    dataset = M3DCAPDataset_eval(
        image_folder='/blue/kgong/wujiong/Experiments/VLM/M3D/Data/data/',
        vqa_file='/blue/kgong/wujiong/Experiments/VLM/M3D/Data/data/M3D_Cap_npy/M3D_Cap.json',
        mode=args.vqa_file
    )

    sampler = torch.utils.data.DistributedSampler(
        dataset,
        num_replicas=world_size,
        rank=rank,
        shuffle=False,
        drop_last=False
    )
    dataloader = torch.utils.data.DataLoader(
        dataset,
        sampler=sampler,
        batch_size=1,
        num_workers=2,
        pin_memory=False,
        collate_fn=lambda x: x[0],
    )
    results = []
    executor = concurrent.futures.ThreadPoolExecutor()
    for item in tqdm.tqdm(dataloader):
        with torch.no_grad():
            assert item['images'] is not None
            result = model.predict_forward(
                video=item['images'],
                text=item['text_prompt'],
                tokenizer=tokenizer,
            )

        text_idx = 0
        text_prediction = result['prediction']
        # if len(result['prediction_masks']) > 0:
        #     mask_prediction = result['prediction_masks'][text_idx]
        # else:
        #     print(text_prediction)
        #     mask_prediction = np.zeros((item['length'], item['ori_height'], item['ori_width']), dtype=np.uint8)
        #
        # dice = dice_score(mask_prediction, item['masks'])

        # if args.submit:
        #     async_func(executor, mask_save, item=item, mask_prediction=mask_prediction, work_dir=work_dir)
        #     encoded_mask = None
        # else:
        #     encoded_mask = mask_to_rle(mask_prediction)

        result = {
            'index': item['index'],
            'video_id': item['video_id'],
            'exp_id': str(item['exp_id']),
            'text_prediction': text_prediction,
            'frames': item['frames'],
            'exp': item['text_prompt'],
            'labels': item['labels']

        }
        results.append(result)

    executor.shutdown(wait=True)
    print(f'[Rank {rank}] : Finished.')

    if not args.submit:
        results = collect_results_cpu(results, len(dataset))
        if get_rank() == 0:
            final_results = {}
            for item in results:
                vid_id = item['video_id']
                exp_id = str(item['exp_id'])
                if exp_id not in final_results:
                    final_results[exp_id] = []
                # assert exp_id not in final_results[vid_id]
                final_results[exp_id].append(item)
            os.makedirs(work_dir, exist_ok=True)
            json.dump(final_results, open(f'{work_dir}/results_closed_vqa.json', 'w'))

    if rank == 0:
        print('Done')