File size: 2,988 Bytes
0dc87cd | 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 | """VLM fallback using Qwen2.5-VL for low-confidence product identification."""
import numpy as np
from PIL import Image
QWEN_MODEL_ID = "Qwen/Qwen2.5-VL-3B-Instruct"
_model = None
_processor = None
def load_vlm(device: str = "cuda"):
global _model, _processor
if _model is not None:
return _model, _processor
from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
import torch
print(f"[vlm] Loading Qwen2.5-VL ({QWEN_MODEL_ID}) on {device}...")
_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
QWEN_MODEL_ID,
torch_dtype=torch.float16,
device_map=device,
)
_processor = AutoProcessor.from_pretrained(QWEN_MODEL_ID)
_model.eval()
print("[vlm] Qwen2.5-VL loaded.")
return _model, _processor
def identify(crop_bgr: np.ndarray, product_names: list[str], device: str = "cuda") -> str:
"""Use Qwen2.5-VL to identify a product crop from a list of known products.
Returns the matched product name or 'unknown'.
"""
import torch
from qwen_vl_utils import process_vision_info
model, processor = load_vlm(device)
rgb = crop_bgr[:, :, ::-1]
pil_img = Image.fromarray(rgb)
product_list = "\n".join(f"- {name}" for name in product_names)
messages = [
{
"role": "user",
"content": [
{"type": "image", "image": pil_img},
{"type": "text", "text": (
"This is a product on a retail shelf. Identify which product it is "
"from this list:\n"
f"{product_list}\n\n"
"If it matches one, reply with ONLY the exact product name from the list. "
"If it does not match any, reply with ONLY the word 'unknown'."
)},
],
}
]
text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
image_inputs, video_inputs = process_vision_info(messages)
inputs = processor(
text=[text],
images=image_inputs,
videos=video_inputs,
padding=True,
return_tensors="pt",
).to(device)
with torch.no_grad():
output_ids = model.generate(**inputs, max_new_tokens=50)
input_len = inputs.input_ids.shape[1]
response = processor.batch_decode(output_ids[:, input_len:], skip_special_tokens=True)[0].strip()
for name in product_names:
if name.lower() in response.lower():
return name
return response
if __name__ == "__main__":
import cv2
import sys
img_path = sys.argv[1] if len(sys.argv) > 1 else "outputs/crops/shelf_03_crop_000.jpg"
crop = cv2.imread(img_path)
product_names = [
"Parle-G Gold Biscuits", "Britannia Good Day Butter Cookies",
"Oreo Chocolate Creme Biscuits", "Sunfeast Dark Fantasy Choco Fills",
]
result = identify(crop, product_names)
print(f"VLM result: {result}")
|