336 lines
12 KiB
Python
336 lines
12 KiB
Python
import base64
|
|
import io
|
|
import json
|
|
import os
|
|
import re
|
|
import threading
|
|
from pathlib import Path
|
|
|
|
from flask import Flask, jsonify, request
|
|
from PIL import Image
|
|
import torch
|
|
from torchvision import transforms
|
|
from transformers import AutoModel, AutoTokenizer
|
|
|
|
|
|
app = Flask(__name__)
|
|
app.config["MAX_CONTENT_LENGTH"] = 50 * 1024 * 1024
|
|
|
|
MODEL_PATH = Path(os.getenv("INTERNVL3_MODEL_PATH", r"D:\go-demo\data\InternVL3-8B"))
|
|
INPUT_SIZE = int(os.getenv("INTERNVL3_IMAGE_SIZE", "448"))
|
|
MAX_TILES = int(os.getenv("INTERNVL3_MAX_TILES", "4"))
|
|
MAX_NEW_TOKENS = int(os.getenv("INTERNVL3_MAX_NEW_TOKENS", "512"))
|
|
TEMPERATURE = float(os.getenv("INTERNVL3_TEMPERATURE", "0.0"))
|
|
USE_FLASH_ATTN = os.getenv("INTERNVL3_USE_FLASH_ATTN", "").lower() in {"1", "true", "yes"}
|
|
LOAD_IN_8BIT = os.getenv("INTERNVL3_LOAD_IN_8BIT", "").lower() in {"1", "true", "yes"}
|
|
LOAD_IN_4BIT = os.getenv("INTERNVL3_LOAD_IN_4BIT", "").lower() in {"1", "true", "yes"}
|
|
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
_lock = threading.Lock()
|
|
_model = None
|
|
_tokenizer = None
|
|
_ready = False
|
|
_load_error = None
|
|
|
|
IMAGENET_MEAN = (0.485, 0.456, 0.406)
|
|
IMAGENET_STD = (0.229, 0.224, 0.225)
|
|
JSON_RE = re.compile(r"\{.*\}", re.S)
|
|
|
|
|
|
def load_model():
|
|
global _model, _tokenizer, _ready, _load_error
|
|
with _lock:
|
|
if _ready:
|
|
return _model, _tokenizer
|
|
if _load_error is not None:
|
|
raise RuntimeError(_load_error)
|
|
if not MODEL_PATH.exists():
|
|
_load_error = f"model path not found: {MODEL_PATH}"
|
|
raise RuntimeError(_load_error)
|
|
|
|
kwargs = {
|
|
"trust_remote_code": True,
|
|
"low_cpu_mem_usage": True,
|
|
}
|
|
if USE_FLASH_ATTN:
|
|
kwargs["use_flash_attn"] = True
|
|
if LOAD_IN_8BIT or LOAD_IN_4BIT:
|
|
try:
|
|
from transformers import BitsAndBytesConfig
|
|
|
|
kwargs["quantization_config"] = BitsAndBytesConfig(
|
|
load_in_8bit=LOAD_IN_8BIT,
|
|
load_in_4bit=LOAD_IN_4BIT,
|
|
)
|
|
except Exception as exc:
|
|
_load_error = f"bitsandbytes quantization unavailable: {exc}"
|
|
raise RuntimeError(_load_error)
|
|
|
|
if DEVICE == "cuda" and not (LOAD_IN_8BIT or LOAD_IN_4BIT):
|
|
kwargs["torch_dtype"] = torch.bfloat16
|
|
kwargs["device_map"] = "auto"
|
|
else:
|
|
kwargs["torch_dtype"] = torch.float32 if DEVICE == "cpu" else torch.bfloat16
|
|
|
|
app.logger.info("Loading InternVL3 from %s on %s", MODEL_PATH, DEVICE)
|
|
_tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True, use_fast=False)
|
|
_model = AutoModel.from_pretrained(MODEL_PATH, **kwargs).eval()
|
|
_ready = True
|
|
app.logger.info("InternVL3 ready")
|
|
return _model, _tokenizer
|
|
|
|
|
|
def build_transform():
|
|
return transforms.Compose(
|
|
[
|
|
transforms.Resize((INPUT_SIZE, INPUT_SIZE), interpolation=transforms.InterpolationMode.BICUBIC),
|
|
transforms.ToTensor(),
|
|
transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
|
|
]
|
|
)
|
|
|
|
|
|
def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
|
|
best_ratio_diff = float("inf")
|
|
best_ratio = (1, 1)
|
|
area = width * height
|
|
for ratio in target_ratios:
|
|
target_aspect_ratio = ratio[0] / ratio[1]
|
|
ratio_diff = abs(aspect_ratio - target_aspect_ratio)
|
|
if ratio_diff < best_ratio_diff:
|
|
best_ratio_diff = ratio_diff
|
|
best_ratio = ratio
|
|
elif ratio_diff == best_ratio_diff:
|
|
if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
|
|
best_ratio = ratio
|
|
return best_ratio
|
|
|
|
|
|
def dynamic_preprocess(image, min_num=1, max_num=12, image_size=448, use_thumbnail=False):
|
|
orig_width, orig_height = image.size
|
|
aspect_ratio = orig_width / orig_height
|
|
|
|
target_ratios = set(
|
|
(i, j)
|
|
for n in range(min_num, max_num + 1)
|
|
for i in range(1, n + 1)
|
|
for j in range(1, n + 1)
|
|
if i * j <= max_num and i * j >= min_num
|
|
)
|
|
target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
|
|
target_aspect_ratio = find_closest_aspect_ratio(
|
|
aspect_ratio, target_ratios, orig_width, orig_height, image_size
|
|
)
|
|
|
|
target_width = image_size * target_aspect_ratio[0]
|
|
target_height = image_size * target_aspect_ratio[1]
|
|
blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
|
|
resized_img = image.resize((target_width, target_height), Image.Resampling.BICUBIC)
|
|
processed_images = []
|
|
for i in range(blocks):
|
|
box = (
|
|
(i % (target_width // image_size)) * image_size,
|
|
(i // (target_width // image_size)) * image_size,
|
|
((i % (target_width // image_size)) + 1) * image_size,
|
|
((i // (target_width // image_size)) + 1) * image_size,
|
|
)
|
|
processed_images.append(resized_img.crop(box))
|
|
if use_thumbnail and len(processed_images) != 1:
|
|
processed_images.append(image.resize((image_size, image_size), Image.Resampling.BICUBIC))
|
|
return processed_images
|
|
|
|
|
|
def split_tiles(image):
|
|
return dynamic_preprocess(image, max_num=MAX_TILES, image_size=INPUT_SIZE, use_thumbnail=True)
|
|
|
|
|
|
def image_to_tensor(image):
|
|
return build_transform()(image.convert("RGB"))
|
|
|
|
|
|
def decode_image(payload):
|
|
if payload.get("image_path"):
|
|
return Image.open(payload["image_path"]).convert("RGB")
|
|
raw = payload.get("image_base64")
|
|
if not raw:
|
|
raise ValueError("image_base64 or image_path is required")
|
|
data = base64.b64decode(raw)
|
|
return Image.open(io.BytesIO(data)).convert("RGB")
|
|
|
|
|
|
def prompt_for_product(product):
|
|
sku = product.get("sku", "")
|
|
brand = product.get("brand", "")
|
|
category = product.get("category", "")
|
|
material = product.get("material", "")
|
|
style = product.get("style_name", "")
|
|
color = product.get("color_tone", "")
|
|
finish = product.get("finish", "")
|
|
return f"""<image>
|
|
You are a flooring material analysis engine.
|
|
Analyze only the floor material in the image.
|
|
Return valid JSON only. No markdown. No commentary.
|
|
|
|
Schema:
|
|
{{
|
|
"description": "",
|
|
"material_type": "",
|
|
"surface_finish": "",
|
|
"grain_type": "",
|
|
"color_family": "",
|
|
"visual_style": "",
|
|
"variation": "",
|
|
"gloss_level": "",
|
|
"tags": [],
|
|
"confidence": 0.0,
|
|
"field_confidence": {{}}
|
|
}}
|
|
|
|
Rules:
|
|
- material_type should be a short normalized label like wood, tile, stone, vinyl, laminate, spc, lvp, wpc, ceramic, porcelain, or generic.
|
|
- gloss_level must be High, Medium, or Low.
|
|
- confidence and field_confidence values must be numbers from 0 to 1.
|
|
- tags should contain 5 to 12 short material keywords.
|
|
- Focus on flooring surface only, not room decor.
|
|
|
|
Product metadata:
|
|
SKU: {sku}
|
|
Brand: {brand}
|
|
Category: {category}
|
|
Material: {material}
|
|
Style: {style}
|
|
Color tone: {color}
|
|
Finish: {finish}
|
|
"""
|
|
|
|
|
|
def extract_json_text(text):
|
|
if not text:
|
|
raise ValueError("empty model output")
|
|
match = JSON_RE.search(text)
|
|
if not match:
|
|
raise ValueError("no JSON object found in model output")
|
|
return match.group(0)
|
|
|
|
|
|
def normalize_output(data, product):
|
|
if not isinstance(data, dict):
|
|
raise ValueError("semantic output must be a JSON object")
|
|
|
|
def titleish(value):
|
|
if not isinstance(value, str):
|
|
return ""
|
|
return value.strip()
|
|
|
|
semantic = {
|
|
"provider": "internvl3",
|
|
"model": str(MODEL_PATH.name),
|
|
"description": titleish(data.get("description")),
|
|
"material_type": titleish(data.get("material_type")).lower(),
|
|
"surface_finish": titleish(data.get("surface_finish")),
|
|
"grain_type": titleish(data.get("grain_type")),
|
|
"color_family": titleish(data.get("color_family")),
|
|
"visual_style": titleish(data.get("visual_style")),
|
|
"variation": titleish(data.get("variation")),
|
|
"gloss_level": titleish(data.get("gloss_level")).title(),
|
|
"tags": data.get("tags") or [],
|
|
"confidence": float(data.get("confidence") or 0.0),
|
|
"field_confidence": data.get("field_confidence") or {},
|
|
}
|
|
if not semantic["description"]:
|
|
semantic["description"] = f"{product.get('material', '')} flooring"
|
|
if not semantic["material_type"]:
|
|
semantic["material_type"] = "generic"
|
|
if semantic["gloss_level"] not in {"High", "Medium", "Low"}:
|
|
semantic["gloss_level"] = "Medium"
|
|
if not isinstance(semantic["tags"], list):
|
|
semantic["tags"] = []
|
|
semantic["tags"] = [str(tag).strip() for tag in semantic["tags"] if str(tag).strip()]
|
|
if len(semantic["tags"]) < 5:
|
|
semantic["tags"].extend(
|
|
[product.get("category", ""), product.get("material", ""), product.get("brand", "")]
|
|
)
|
|
semantic["tags"] = [tag for tag in semantic["tags"] if tag]
|
|
if not isinstance(semantic["field_confidence"], dict):
|
|
semantic["field_confidence"] = {}
|
|
return semantic
|
|
|
|
|
|
def analyze_semantic(image, product):
|
|
model, tokenizer = load_model()
|
|
prompt = prompt_for_product(product)
|
|
pixel_values = torch.stack([image_to_tensor(tile) for tile in split_tiles(image)])
|
|
input_device = next(model.parameters()).device
|
|
pixel_values = pixel_values.to(input_device)
|
|
generation_config = {
|
|
"max_new_tokens": MAX_NEW_TOKENS,
|
|
"do_sample": False,
|
|
"temperature": TEMPERATURE,
|
|
}
|
|
question = prompt
|
|
last_error = None
|
|
for attempt in range(3):
|
|
try:
|
|
with torch.no_grad():
|
|
response = model.chat(tokenizer, pixel_values, question, generation_config)
|
|
text = response[0] if isinstance(response, (list, tuple)) else response
|
|
semantic = normalize_output(json.loads(extract_json_text(text)), product)
|
|
semantic["confidence"] = max(0.0, min(1.0, float(semantic["confidence"])))
|
|
if not semantic["field_confidence"]:
|
|
semantic["field_confidence"] = {
|
|
"description": semantic["confidence"],
|
|
"material_type": semantic["confidence"],
|
|
"surface_finish": semantic["confidence"],
|
|
"grain_type": semantic["confidence"],
|
|
"color_family": semantic["confidence"],
|
|
"visual_style": semantic["confidence"],
|
|
"variation": semantic["confidence"],
|
|
"gloss_level": semantic["confidence"],
|
|
}
|
|
return semantic
|
|
except Exception as exc:
|
|
last_error = exc
|
|
question = (
|
|
prompt
|
|
+ "\nYour previous output was invalid. Return only a JSON object matching the schema."
|
|
+ f"\nError: {exc}"
|
|
)
|
|
raise RuntimeError(f"semantic analysis failed after retries: {last_error}")
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
try:
|
|
model, _ = load_model()
|
|
device = str(next(model.parameters()).device)
|
|
return jsonify({
|
|
"status": "ok",
|
|
"ready": True,
|
|
"model_path": str(MODEL_PATH),
|
|
"device": device,
|
|
"input_size": INPUT_SIZE,
|
|
"max_tiles": MAX_TILES,
|
|
})
|
|
except Exception as exc:
|
|
return jsonify({
|
|
"status": "error",
|
|
"ready": False,
|
|
"model_path": str(MODEL_PATH),
|
|
"error": str(exc),
|
|
}), 500
|
|
|
|
|
|
@app.post("/semantic")
|
|
def semantic():
|
|
payload = request.get_json(force=True, silent=False)
|
|
image = decode_image(payload)
|
|
product = payload.get("product") or {}
|
|
semantic = analyze_semantic(image, product)
|
|
return jsonify(semantic)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
port = int(os.getenv("INTERNVL3_PORT", "5300"))
|
|
app.logger.info("Starting InternVL3 vision server on :%s", port)
|
|
app.run(host="0.0.0.0", port=port, debug=False)
|