| import argparse |
| import copy |
| import math |
| import os |
| import torch |
| import tqdm |
| from pycocotools import mask as _mask |
| import numpy as np |
| import random |
|
|
| from transformers import (AutoModel, AutoModelForCausalLM, AutoTokenizer, |
| BitsAndBytesConfig, CLIPImageProcessor, |
| CLIPVisionModel, GenerationConfig) |
|
|
| from utils import _init_dist_pytorch, get_dist_info, get_rank, collect_results_cpu |
| from datasets import RESDataset |
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser(description='RefCocoSeg') |
| parser.add_argument('model_path', help='hf model path.') |
| parser.add_argument( |
| '--dataset', |
| choices=DATASETS_ATTRIBUTES.keys(), |
| default='refcoco', |
| help='Specify a ref dataset') |
| parser.add_argument( |
| '--split', |
| default='val', |
| help='Specify a split') |
| 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) |
| args = parser.parse_args() |
| if 'LOCAL_RANK' not in os.environ: |
| os.environ['LOCAL_RANK'] = str(args.local_rank) |
| return args |
|
|
| DATASETS_ATTRIBUTES = { |
| 'refcoco': {'splitBy': "unc", 'dataset_name': 'refcoco'}, |
| 'refcoco_plus': {'splitBy': "unc", 'dataset_name': 'refcoco_plus'}, |
| 'refcocog': {'splitBy': "umd", 'dataset_name': 'refcocog'}, |
| } |
|
|
| IMAGE_FOLDER = './data/glamm_data/images/coco2014/train2014/' |
| DATA_PATH = './data/ref_seg/' |
|
|
|
|
| def main(): |
| args = parse_args() |
|
|
| if args.launcher != 'none': |
| _init_dist_pytorch('nccl') |
| rank, world_size = get_dist_info() |
| torch.cuda.set_device(rank) |
| else: |
| rank = 0 |
| world_size = 1 |
|
|
| |
| model = AutoModel.from_pretrained( |
| args.model_path, |
| torch_dtype=torch.bfloat16, |
| low_cpu_mem_usage=True, |
| use_flash_attn=True, |
| trust_remote_code=True, |
| ).eval().cuda() |
|
|
| tokenizer = AutoTokenizer.from_pretrained( |
| args.model_path, |
| trust_remote_code=True, |
| ) |
| dataset_info = DATASETS_ATTRIBUTES[args.dataset] |
|
|
| dataset = RESDataset( |
| image_folder=IMAGE_FOLDER, |
| dataset_name=dataset_info['dataset_name'], |
| data_path=DATA_PATH, |
| split=args.split, |
| ) |
|
|
| results = [] |
| n_samples = len(dataset) |
| per_rank_samples = math.ceil(n_samples / world_size) + 1 |
| per_rank_ids = range(per_rank_samples * rank, |
| min(n_samples, per_rank_samples * (rank + 1))) |
| for idx in tqdm.tqdm(per_rank_ids): |
| data_batch = dataset[idx] |
| prediction = {'img_id': data_batch['img_id'], 'gt_masks': data_batch['gt_masks']} |
| prediction['gt_masks'] = mask_to_rle(prediction['gt_masks'].cpu().numpy()) |
| del data_batch['img_id'], data_batch['gt_masks'] |
|
|
| texts = data_batch['text'] |
| del data_batch['text'] |
| pred_masks = [] |
| for text in texts: |
| _data_batch = copy.deepcopy(data_batch) |
| _data_batch['text'] = text |
| pred_mask = model.predict_forward(**_data_batch, tokenizer=tokenizer)['prediction_masks'] |
| if len(pred_mask) == 0: |
| |
| print("No seg pred !!!") |
| pred_masks.append(None) |
| else: |
| _ret_mask = pred_mask[0].cpu().numpy() |
| _ret_mask = mask_to_rle(_ret_mask) |
| pred_masks.append(_ret_mask) |
|
|
| prediction.update({'prediction_masks': pred_masks}) |
| results.append(prediction) |
|
|
| tmpdir = './dist_test_temp_res_' + args.dataset + args.split + args.model_path.replace('/', '').replace('.', '') |
| results = collect_results_cpu(results, len(dataset), tmpdir=tmpdir) |
| if get_rank() == 0: |
| metric = dataset.evaluate(results, './work_dirs') |
| print(metric) |
|
|
| def mask_to_rle(mask): |
| rle = [] |
| for m in mask: |
| rle.append(_mask.encode(np.asfortranarray(m.astype(np.uint8)))) |
| rle[-1]['counts'] = rle[-1]['counts'].decode() |
| return rle |
|
|
| if __name__ == '__main__': |
| main() |
|
|
