Spaces:
Paused
Paused
| import argparse | |
| import os | |
| import sys | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from transformers import AutoTokenizer, BitsAndBytesConfig, CLIPImageProcessor | |
| from model.LISA import LISAForCausalLM | |
| from model.llava import conversation as conversation_lib | |
| from model.llava.mm_utils import tokenizer_image_token | |
| from model.segment_anything.utils.transforms import ResizeLongestSide | |
| from utils.utils import (DEFAULT_IM_END_TOKEN, DEFAULT_IM_START_TOKEN, | |
| DEFAULT_IMAGE_TOKEN, IMAGE_TOKEN_INDEX) | |
| def parse_args(args): | |
| parser = argparse.ArgumentParser(description="LISA chat") | |
| parser.add_argument("--version", default="xinlai/LISA-13B-llama2-v1") | |
| parser.add_argument("--vis_save_path", default="./vis_output", type=str) | |
| parser.add_argument( | |
| "--precision", | |
| default="bf16", | |
| type=str, | |
| choices=["fp32", "bf16", "fp16"], | |
| help="precision for inference", | |
| ) | |
| parser.add_argument("--image_size", default=1024, type=int, help="image size") | |
| parser.add_argument("--model_max_length", default=512, type=int) | |
| parser.add_argument("--lora_r", default=8, type=int) | |
| parser.add_argument( | |
| "--vision-tower", default="openai/clip-vit-large-patch14", type=str | |
| ) | |
| parser.add_argument("--local-rank", default=0, type=int, help="node rank") | |
| parser.add_argument("--load_in_8bit", action="store_true", default=False) | |
| parser.add_argument("--load_in_4bit", action="store_true", default=False) | |
| parser.add_argument("--use_mm_start_end", action="store_true", default=True) | |
| parser.add_argument( | |
| "--conv_type", | |
| default="llava_v1", | |
| type=str, | |
| choices=["llava_v1", "llava_llama_2"], | |
| ) | |
| return parser.parse_args(args) | |
| def preprocess( | |
| x, | |
| pixel_mean=torch.Tensor([123.675, 116.28, 103.53]).view(-1, 1, 1), | |
| pixel_std=torch.Tensor([58.395, 57.12, 57.375]).view(-1, 1, 1), | |
| img_size=1024, | |
| ) -> torch.Tensor: | |
| """Normalize pixel values and pad to a square input.""" | |
| # Normalize colors | |
| x = (x - pixel_mean) / pixel_std | |
| # Pad | |
| h, w = x.shape[-2:] | |
| padh = img_size - h | |
| padw = img_size - w | |
| x = F.pad(x, (0, padw, 0, padh)) | |
| return x | |
| def main(args): | |
| args = parse_args(args) | |
| os.makedirs(args.vis_save_path, exist_ok=True) | |
| # Create model | |
| tokenizer = AutoTokenizer.from_pretrained( | |
| args.version, | |
| cache_dir=None, | |
| model_max_length=args.model_max_length, | |
| padding_side="right", | |
| use_fast=False, | |
| ) | |
| tokenizer.pad_token = tokenizer.unk_token | |
| args.seg_token_idx = tokenizer("[SEG]", add_special_tokens=False).input_ids[0] | |
| torch_dtype = torch.float32 | |
| if args.precision == "bf16": | |
| torch_dtype = torch.bfloat16 | |
| elif args.precision == "fp16": | |
| torch_dtype = torch.half | |
| kwargs = {"torch_dtype": torch_dtype} | |
| if args.load_in_4bit: | |
| kwargs.update( | |
| { | |
| "torch_dtype": torch.half, | |
| "device_map": "auto", | |
| "load_in_4bit": True, | |
| "quantization_config": BitsAndBytesConfig( | |
| load_in_4bit=True, | |
| bnb_4bit_compute_dtype=torch.float16, | |
| bnb_4bit_use_double_quant=True, | |
| bnb_4bit_quant_type="nf4", | |
| llm_int8_skip_modules=["visual_model"], | |
| ), | |
| } | |
| ) | |
| elif args.load_in_8bit: | |
| kwargs.update( | |
| { | |
| "torch_dtype": torch.half, | |
| "device_map": "auto", | |
| "quantization_config": BitsAndBytesConfig( | |
| llm_int8_skip_modules=["visual_model"], | |
| load_in_8bit=True, | |
| ), | |
| } | |
| ) | |
| model = LISAForCausalLM.from_pretrained( | |
| args.version, low_cpu_mem_usage=True, seg_token_idx=args.seg_token_idx, **kwargs | |
| ) | |
| model.config.eos_token_id = tokenizer.eos_token_id | |
| model.config.bos_token_id = tokenizer.bos_token_id | |
| model.config.pad_token_id = tokenizer.pad_token_id | |
| model.get_model().initialize_vision_modules(model.get_model().config) | |
| vision_tower = model.get_model().get_vision_tower() | |
| vision_tower.to(dtype=torch_dtype) | |
| if args.precision == "bf16": | |
| model = model.bfloat16().cuda() | |
| elif ( | |
| args.precision == "fp16" and (not args.load_in_4bit) and (not args.load_in_8bit) | |
| ): | |
| vision_tower = model.get_model().get_vision_tower() | |
| model.model.vision_tower = None | |
| import deepspeed | |
| model_engine = deepspeed.init_inference( | |
| model=model, | |
| dtype=torch.half, | |
| replace_with_kernel_inject=True, | |
| replace_method="auto", | |
| ) | |
| model = model_engine.module | |
| model.model.vision_tower = vision_tower.half().cuda() | |
| elif args.precision == "fp32": | |
| model = model.float().cuda() | |
| vision_tower = model.get_model().get_vision_tower() | |
| vision_tower.to(device=args.local_rank) | |
| clip_image_processor = CLIPImageProcessor.from_pretrained(model.config.vision_tower) | |
| transform = ResizeLongestSide(args.image_size) | |
| model.eval() | |
| while True: | |
| conv = conversation_lib.conv_templates[args.conv_type].copy() | |
| conv.messages = [] | |
| prompt = input("Please input your prompt: ") | |
| prompt = DEFAULT_IMAGE_TOKEN + "\n" + prompt | |
| if args.use_mm_start_end: | |
| replace_token = ( | |
| DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN | |
| ) | |
| prompt = prompt.replace(DEFAULT_IMAGE_TOKEN, replace_token) | |
| conv.append_message(conv.roles[0], prompt) | |
| conv.append_message(conv.roles[1], "") | |
| prompt = conv.get_prompt() | |
| image_path = input("Please input the image path: ") | |
| if not os.path.exists(image_path): | |
| print("File not found in {}".format(image_path)) | |
| continue | |
| image_np = cv2.imread(image_path) | |
| image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB) | |
| original_size_list = [image_np.shape[:2]] | |
| image_clip = ( | |
| clip_image_processor.preprocess(image_np, return_tensors="pt")[ | |
| "pixel_values" | |
| ][0] | |
| .unsqueeze(0) | |
| .cuda() | |
| ) | |
| if args.precision == "bf16": | |
| image_clip = image_clip.bfloat16() | |
| elif args.precision == "fp16": | |
| image_clip = image_clip.half() | |
| else: | |
| image_clip = image_clip.float() | |
| image = transform.apply_image(image_np) | |
| resize_list = [image.shape[:2]] | |
| image = ( | |
| preprocess(torch.from_numpy(image).permute(2, 0, 1).contiguous()) | |
| .unsqueeze(0) | |
| .cuda() | |
| ) | |
| if args.precision == "bf16": | |
| image = image.bfloat16() | |
| elif args.precision == "fp16": | |
| image = image.half() | |
| else: | |
| image = image.float() | |
| input_ids = tokenizer_image_token(prompt, tokenizer, return_tensors="pt") | |
| input_ids = input_ids.unsqueeze(0).cuda() | |
| output_ids, pred_masks = model.evaluate( | |
| image_clip, | |
| image, | |
| input_ids, | |
| resize_list, | |
| original_size_list, | |
| max_new_tokens=512, | |
| tokenizer=tokenizer, | |
| ) | |
| output_ids = output_ids[0][output_ids[0] != IMAGE_TOKEN_INDEX] | |
| text_output = tokenizer.decode(output_ids, skip_special_tokens=False) | |
| text_output = text_output.replace("\n", "").replace(" ", " ") | |
| for i, pred_mask in enumerate(pred_masks): | |
| if pred_mask.shape[0] == 0: | |
| continue | |
| pred_mask = pred_mask.detach().cpu().numpy()[0] | |
| pred_mask = pred_mask > 0 | |
| save_path = "{}/{}_mask_{}.jpg".format( | |
| args.vis_save_path, image_path.split("/")[-1].split(".")[0], i | |
| ) | |
| cv2.imwrite(save_path, pred_mask * 100) | |
| print("{} has been saved.".format(save_path)) | |
| save_path = "{}/{}_masked_img_{}.jpg".format( | |
| args.vis_save_path, image_path.split("/")[-1].split(".")[0], i | |
| ) | |
| save_img = image_np.copy() | |
| save_img[pred_mask] = ( | |
| image_np * 0.5 | |
| + pred_mask[:, :, None].astype(np.uint8) * np.array([255, 0, 0]) * 0.5 | |
| )[pred_mask] | |
| save_img = cv2.cvtColor(save_img, cv2.COLOR_RGB2BGR) | |
| cv2.imwrite(save_path, save_img) | |
| print("{} has been saved.".format(save_path)) | |
| if __name__ == "__main__": | |
| main(sys.argv[1:]) | |