Spaces:
Running
on
Zero
Running
on
Zero
File size: 5,097 Bytes
46861c5 |
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 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 |
# --------------------------------------------------------
# Copyright (2025) Bytedance Ltd. and/or its affiliates
# Licensed under the Apache License, Version 2.0 (the "License")
# Grasp Any Region Project
# Written by Haochen Wang
# --------------------------------------------------------
import argparse
import json
import os
import numpy as np
import torch
from PIL import Image
from pycocotools import mask as mask_utils
from pycocotools.coco import COCO
from tqdm import tqdm
from transformers import AutoModel, AutoProcessor, GenerationConfig
from evaluation.eval_dataset import SingleRegionCaptionDataset
TORCH_DTYPE_MAP = dict(fp16=torch.float16, bf16=torch.bfloat16, fp32=torch.float32)
def parse_args():
parser = argparse.ArgumentParser(
description="Inference of Grasp Any Region models on DLC-Bench."
)
parser.add_argument(
"--model_name_or_path",
help="HF model name or path",
default="HaochenWang/GAR-1B",
)
parser.add_argument(
"--cache_name",
help="cache name to save model outputs.",
default="gar_1b",
)
parser.add_argument(
"--data_type",
help="data dtype",
type=str,
choices=["fp16", "bf16", "fp32"],
default="bf16",
)
parser.add_argument(
"--anno_file",
help="path to the annotation file.",
default="evaluation/DLC-Bench/annotations/annotations.json",
)
parser.add_argument(
"--image_folder",
help="the folder of images",
default="evaluation/DLC-Bench/annotations",
)
parser.add_argument(
"--seed",
type=int,
default=0,
help="Random seed for reproducible text generation",
)
args = parser.parse_args()
return args
def select_ann(coco, img_id, area_min=None, area_max=None):
cat_ids = coco.getCatIds()
ann_ids = coco.getAnnIds(imgIds=[img_id], catIds=cat_ids, iscrowd=None)
if area_min is not None:
ann_ids = [
ann_id for ann_id in ann_ids if coco.anns[ann_id]["area"] >= area_min
]
if area_max is not None:
ann_ids = [
ann_id for ann_id in ann_ids if coco.anns[ann_id]["area"] <= area_max
]
return ann_ids
def main():
args = parse_args()
data_dtype = TORCH_DTYPE_MAP[args.data_type]
torch.manual_seed(args.seed)
# init ditribution for dispatch_modules in LLM
torch.cuda.set_device(0)
torch.distributed.init_process_group(backend="nccl")
# build HF model
model = AutoModel.from_pretrained(
args.model_name_or_path,
trust_remote_code=True,
torch_dtype=data_dtype,
)
model.cuda()
model.eval()
processor = AutoProcessor.from_pretrained(
args.model_name_or_path,
trust_remote_code=True,
)
model_outputs = {}
cache_name = args.cache_name
# This coco instance is actually an o365 subset. This is for code reuse.
coco = COCO(args.anno_file)
img_ids = list(coco.imgs.keys())
num_anns = len(coco.anns)
pbar = tqdm(total=num_anns)
for img_id in img_ids:
ann_ids = select_ann(coco, img_id)
img_info = coco.loadImgs(img_id)[0]
for i, ann_id in enumerate(ann_ids):
if ann_id in model_outputs.keys():
pbar.update(1)
continue
anns = coco.loadAnns([ann_id])
mask = coco.annToMask(anns[0])
img_path = os.path.join(args.image_folder, "images", img_info["file_name"])
img = Image.open(img_path)
prompt_number = model.config.prompt_numbers
prompt_tokens = [f"<Prompt{i_p}>" for i_p in range(prompt_number)] + [
"<NO_Prompt>"
]
dataset = SingleRegionCaptionDataset(
image=img,
mask=mask,
processor=processor,
prompt_number=prompt_number,
visual_prompt_tokens=prompt_tokens,
data_dtype=data_dtype,
)
data_sample = dataset[0]
with torch.no_grad():
generate_ids = model.generate(
**data_sample,
generation_config=GenerationConfig(
max_new_tokens=1024,
do_sample=False,
eos_token_id=processor.tokenizer.eos_token_id,
pad_token_id=processor.tokenizer.pad_token_id,
),
return_dict=True,
)
outputs = processor.tokenizer.decode(
generate_ids.sequences[0], skip_special_tokens=True
).strip()
print(outputs) # Print model output for this image
model_outputs[ann_id] = outputs
pbar.update(1)
pbar.close()
with open(f"evaluation/DLC-Bench/model_outputs/{cache_name}.json", "w") as file:
json.dump(model_outputs, file, indent=4, ensure_ascii=False)
print(f"Cache name: {cache_name}")
if __name__ == "__main__":
main()
|