File size: 4,149 Bytes
032e687 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
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
# build model
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:
# give a zero mask
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()
|