FloorMaterialAnalyzer/tools/vision_server/server.py

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)