#!/usr/bin/env python # -*- coding: utf-8 -*- # Flask-based service to receive an image and a prompt, perform vision-language model inference, and return a segmentation mask. from __future__ import absolute_import, print_function, division from flask import Flask, request, jsonify import os import json import base64 import argparse import numpy as np import cv2 import torch import torch.nn.functional as F from transformers import AutoTokenizer, BitsAndBytesConfig, CLIPImageProcessor from PIL import Image from io import BytesIO # Custom model and utility imports from model.AffordanceVLM import AffordanceVLMForCausalLM 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 ) app = Flask(__name__) # --------------------------- # Argument parser for model config # --------------------------- def parse_args(args): parser = argparse.ArgumentParser(description="AffordanceVLM Flask Service") parser.add_argument("--version", default="/data/AffordanceNet/exps/AffordanceVLM-7B") parser.add_argument("--vis_save_path", default="./vis_output/ur5_samples", type=str) parser.add_argument("--precision", default="bf16", choices=["fp32", "bf16", "fp16"]) parser.add_argument("--image_size", default=1024, type=int) 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") parser.add_argument("--local-rank", default=0, type=int) 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", choices=["llava_v1", "llava_llama_2"]) return parser.parse_args(args) # --------------------------- # Model initialization # --------------------------- args = parse_args(None) os.makedirs(args.vis_save_path, exist_ok=True) # Load tokenizer and add custom tokens tokenizer = AutoTokenizer.from_pretrained(args.version, 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] args.aff_token_idx = tokenizer("[AFF]", add_special_tokens=False).input_ids[0] # Set precision torch_dtype = { "bf16": torch.bfloat16, "fp16": torch.half, "fp32": torch.float32 }[args.precision] # Optional quantization configs kwargs = {"torch_dtype": torch_dtype} if args.load_in_4bit: kwargs.update({ "torch_dtype": torch.half, "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, "quantization_config": BitsAndBytesConfig( load_in_8bit=True, llm_int8_skip_modules=["visual_model"], ), }) # Load model model = AffordanceVLMForCausalLM.from_pretrained( args.version, vision_tower=args.vision_tower, seg_token_idx=args.seg_token_idx, aff_token_idx=args.aff_token_idx, low_cpu_mem_usage=True, **kwargs ) # Set special tokens 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 # Initialize vision modules model.get_model().initialize_vision_modules(model.get_model().config) vision_tower = model.get_model().get_vision_tower().to(dtype=torch_dtype) # Model precision setup if args.precision == "bf16": model = model.bfloat16().cuda() elif args.precision == "fp16" and not args.load_in_4bit and not args.load_in_8bit: model.model.vision_tower = None import deepspeed model = deepspeed.init_inference(model=model, dtype=torch.half, replace_with_kernel_inject=True).module model.model.vision_tower = vision_tower.half().cuda() else: model = model.float().cuda() vision_tower.to(device=args.local_rank) clip_image_processor = CLIPImageProcessor.from_pretrained(model.config.vision_tower) transform = ResizeLongestSide(args.image_size) model.eval() # --------------------------- # Image preprocessing function # --------------------------- 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: x = (x - pixel_mean) / pixel_std h, w = x.shape[-2:] x = F.pad(x, (0, img_size - w, 0, img_size - h)) return x # --------------------------- # Segmentation core logic # --------------------------- def segment(image_path, prompt): conv = conversation_lib.conv_templates[args.conv_type].copy() conv.messages = [] prompt = DEFAULT_IMAGE_TOKEN + "\nYou are an embodied robot. " + prompt if args.use_mm_start_end: prompt = prompt.replace(DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN) conv.append_message(conv.roles[0], prompt) conv.append_message(conv.roles[1], "") prompt = conv.get_prompt() if not os.path.exists(image_path): print(f"File not found: {image_path}") return None image_np = cv2.imread(image_path) image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB) original_size_list = [image_np.shape[:2]] # CLIP preprocessing image_clip = clip_image_processor.preprocess(image_np, return_tensors="pt")["pixel_values"][0].unsqueeze(0).cuda() image_clip = image_clip.to(dtype=torch_dtype) # Resize and normalize image = transform.apply_image(image_np) resize_list = [image.shape[:2]] image = preprocess(torch.from_numpy(image).permute(2, 0, 1)).unsqueeze(0).cuda().to(dtype=torch_dtype) # Tokenize prompt input_ids = tokenizer_image_token(prompt, tokenizer, return_tensors="pt").unsqueeze(0).cuda() # Model inference 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).replace("\n", "").replace(" ", " ") print("text_output:", text_output) # Save predicted masks save_mask_path = "" for i, pred_mask in enumerate(pred_masks): if pred_mask.shape[0] == 0: continue pred_mask = pred_mask.detach().cpu().numpy()[0] > 0 save_mask_path = f"{args.vis_save_path}/{os.path.basename(image_path).split('.')[0]}_mask_{i}.jpg" cv2.imwrite(save_mask_path, pred_mask * 100) print(f"Saved: {save_mask_path}") save_img_path = f"{args.vis_save_path}/{os.path.basename(image_path).split('.')[0]}_masked_img_{i}.jpg" save_img = image_np.copy() save_img[pred_mask] = (image_np * 0.5 + pred_mask[:, :, None] * np.array([255, 0, 0]) * 0.5)[pred_mask] save_img = cv2.cvtColor(save_img, cv2.COLOR_RGB2BGR) cv2.imwrite(save_img_path, save_img) print(f"Saved: {save_img_path}") return save_mask_path # --------------------------- # Convert image to base64 # --------------------------- def img2b64(img): _, buffer = cv2.imencode('.bmp', img) return base64.b64encode(buffer).decode() # --------------------------- # HTTP endpoint: /img_mask # --------------------------- @app.route("/img_mask", methods=['POST', 'GET']) def recv_json(): data = json.loads(request.data) prompt = data.get('prompt', 'no_recv') print("Received prompt:", prompt) # Decode base64 image img_data = base64.b64decode(data['img']) img_np = cv2.imdecode(np.frombuffer(img_data, np.uint8), cv2.IMREAD_COLOR) cv2.imwrite(os.path.join(args.vis_save_path, 'img.jpg'), img_np) # Run segmentation save_path = segment(os.path.join(args.vis_save_path, 'img.jpg'), prompt) img = cv2.imread(save_path) pic_str = img2b64(img) return jsonify({'img': pic_str}) # --------------------------- # App entry point # --------------------------- if __name__ == "__main__": app.run(host='0.0.0.0', port=3200)