| |
| """WeDetect-Ref-2B inference runner (macOS MPS/CPU adapted from official infer_wedetect_ref.py). |
| |
| Pipeline: |
| 1. WeDetect-Base-Uni (YOLO-World style) -> 100 proposals |
| 2. WeDetect-Ref-2B (Qwen3-VL grounding) -> score each <object> token in one forward pass |
| 3. top-1 (or thresholded) bbox -> visualization |
| """ |
| import argparse |
| import copy |
| import json |
| import os |
| import sys |
| import time |
|
|
| import torch |
| from PIL import Image, ImageDraw, ImageFont |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from wedetect_ref.models.qwen3vl_referring import Qwen3VLGroundingForConditionalGeneration |
| from wedetect_ref.models.vision_process import process_vision_info |
| from transformers import AutoProcessor |
| from generate_proposal import SimpleYOLOWorldDetector |
|
|
| FONT_CANDIDATES = [ |
| "/System/Library/Fonts/PingFang.ttc", |
| "/System/Library/Fonts/Hiragino Sans GB.ttc", |
| "/System/Library/Fonts/STHeiti Medium.ttc", |
| "/System/Library/Fonts/Supplemental/Arial Unicode.ttf", |
| ] |
|
|
|
|
| def load_font(size=22): |
| for p in FONT_CANDIDATES: |
| if os.path.exists(p): |
| try: |
| return ImageFont.truetype(p, size) |
| except Exception: |
| continue |
| return ImageFont.load_default() |
|
|
|
|
| def draw_results(image, boxes, scores, query, out_path, proposals=None, proposal_path=None): |
| """Draw final prediction (box+score+label) and optionally all proposals.""" |
| img = image.copy() |
| draw = ImageDraw.Draw(img) |
| font = load_font() |
| colors = ["#FF3B30", "#FF9500", "#34C759", "#007AFF", "#AF52DE", "#FF2D55"] |
| if proposals is not None and proposal_path is not None: |
| pimg = image.copy() |
| pdraw = ImageDraw.Draw(pimg) |
| for b in proposals: |
| pdraw.rectangle(((b[0], b[1]), (b[2], b[3])), outline="#00C7BE", width=2) |
| pimg.save(proposal_path) |
| print(f"[vis] proposals ({len(proposals)}) -> {proposal_path}") |
| for i, (b, s) in enumerate(zip(boxes, scores)): |
| c = colors[i % len(colors)] |
| x1, y1, x2, y2 = b |
| draw.rectangle(((x1, y1), (x2, y2)), outline=c, width=5) |
| label = f"{query} {float(s):.2f}" |
| tb = draw.textbbox((0, 0), label, font=font) |
| tw, th = tb[2] - tb[0], tb[3] - tb[1] |
| ly = y1 - th - 10 if y1 - th - 10 > 0 else y1 + 2 |
| draw.rectangle(((x1, ly), (x1 + tw + 10, ly + th + 8)), fill=c) |
| draw.text((x1 + 5, ly + 3), label, fill="white", font=font) |
| img.save(out_path) |
| print(f"[vis] prediction -> {out_path}") |
|
|
|
|
| def remap_uni_keys(checkpoint): |
| keys = list(checkpoint.keys()) |
| for key in keys: |
| if "backbone" in key: |
| new_key = key.replace("backbone.image_model.model.", "backbone.") |
| checkpoint[new_key] = checkpoint.pop(key) |
| keys = list(checkpoint.keys()) |
| for key in keys: |
| if "bbox_head" in key: |
| new_key = key.replace("bbox_head.head_module.", "bbox_head.") |
| new_key = new_key.replace("0.2.", "0.6.") |
| new_key = new_key.replace("1.2.", "1.6.") |
| new_key = new_key.replace("2.2.", "2.6.") |
| new_key = new_key.replace("1.bn", "4") |
| new_key = new_key.replace("1.conv", "3") |
| new_key = new_key.replace("0.bn", "1") |
| new_key = new_key.replace("0.conv", "0") |
| checkpoint[new_key] = checkpoint.pop(key) |
| return checkpoint |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--wedetect_ref_checkpoint", type=str, default="../models/wedetect-ref-2b") |
| parser.add_argument("--wedetect_uni_checkpoint", type=str, default="../models/wedetect-base-uni/wedetect_base_uni.pth") |
| parser.add_argument("--image", type=str, default="../assets/demo.jpeg") |
| parser.add_argument("--query", type=str, default="a photo of trees and a river") |
| parser.add_argument("--score_thre", type=float, default=-1.0) |
| parser.add_argument("--device", type=str, default="auto", choices=["auto", "mps", "cpu"]) |
| parser.add_argument("--dtype", type=str, default="auto", choices=["auto", "bf16", "fp32"]) |
| parser.add_argument("--out_dir", type=str, default="../outputs") |
| args = parser.parse_args() |
|
|
| device = args.device |
| if device == "auto": |
| device = "mps" if torch.backends.mps.is_available() else "cpu" |
| if args.dtype == "auto": |
| dtype = torch.bfloat16 if device == "mps" else torch.float32 |
| elif args.dtype == "bf16": |
| dtype = torch.bfloat16 |
| else: |
| dtype = torch.float32 |
| print(f"[env] device={device} dtype={dtype}") |
| os.makedirs(args.out_dir, exist_ok=True) |
| stem = os.path.splitext(os.path.basename(args.image))[0] |
|
|
| |
| t0 = time.time() |
| det_model = SimpleYOLOWorldDetector(backbone_size="base", prompt_dim=768, num_prompts=256, num_proposals=100) |
| checkpoint = torch.load(args.wedetect_uni_checkpoint, map_location="cpu", weights_only=True) |
| checkpoint = remap_uni_keys(checkpoint) |
| det_model = det_model.to(device).eval() |
| msg = det_model.load_state_dict(checkpoint, strict=False) |
| n_missing = len(msg.missing_keys) |
| n_unexpected = len(msg.unexpected_keys) |
| print(f"[uni] loaded, missing={n_missing} unexpected={n_unexpected} ({time.time()-t0:.1f}s)") |
| if n_missing > 0: |
| print(" missing sample:", msg.missing_keys[:5]) |
|
|
| with torch.no_grad(): |
| det_out = det_model([args.image]) |
| proposals_np = det_out[0]["bboxes"].float().cpu().numpy() |
| print(f"[uni] proposals: {proposals_np.shape}, took {time.time()-t0:.1f}s") |
|
|
| |
| t1 = time.time() |
| model_kwargs = dict(torch_dtype=dtype, attn_implementation="sdpa") |
| model = Qwen3VLGroundingForConditionalGeneration.from_pretrained(args.wedetect_ref_checkpoint, **model_kwargs) |
| processor = AutoProcessor.from_pretrained(args.wedetect_ref_checkpoint) |
| object_token_index = processor.tokenizer.convert_tokens_to_ids("<object>") |
| model.model.object_token_id = object_token_index |
| model = model.to(device).eval() |
| print(f"[ref] model loaded on {device} ({time.time()-t1:.1f}s), object_token_id={object_token_index}") |
|
|
| image = Image.open(args.image).convert("RGB") |
| ori_shape = [image.size] |
| proposals = [torch.tensor(proposals_np.tolist()).to(device).to(model.dtype)] |
| num_proposals = len(proposals[0]) |
| proposal_str = "<object>" * num_proposals |
|
|
| messages = [ |
| {"role": "user", "content": [{"type": "text", "text": 'Please detect the "%s" in the image' % args.query}]}, |
| {"role": "assistant", "content": [{"type": "text", "text": proposal_str}]}, |
| ] |
| messages[0]["content"] = [{"type": "image", "image": copy.deepcopy(image)}] + messages[0]["content"] |
| image_inputs, video_inputs = process_vision_info(messages, image_patch_size=16) |
| texts = [processor.apply_chat_template(messages, tokenize=False)] |
| model_inputs = processor( |
| text=texts, images=image_inputs, videos=video_inputs, |
| return_tensors="pt", padding=True, do_resize=False, |
| ) |
| model_inputs = model_inputs.to(model.device) |
|
|
| t2 = time.time() |
| with torch.inference_mode(): |
| pred = model( |
| **model_inputs, |
| bboxes=copy.deepcopy(proposals), |
| ori_shapes=ori_shape, |
| bboxes_id=object_token_index, |
| ) |
| print(f"[ref] forward pass took {time.time()-t2:.1f}s, seq_len={model_inputs['input_ids'].shape}") |
|
|
| proposal_positions = model_inputs["input_ids"] == object_token_index |
| pred_scores_all = pred.logits.sigmoid()[proposal_positions].view(-1) |
| pred_bboxes = proposals[0].clone().float() |
|
|
| if args.score_thre < 0: |
| topk_values, topk_indexes = torch.topk(pred_scores_all.view(-1), 1, dim=0) |
| pred_bboxes = pred_bboxes[topk_indexes] |
| pred_scores = topk_values |
| else: |
| mask = pred_scores_all > args.score_thre |
| pred_bboxes = pred_bboxes[mask] |
| pred_scores = pred_scores_all[mask] |
|
|
| boxes_xyxy = pred_bboxes.cpu().tolist() |
| scores = pred_scores.cpu().tolist() |
| print(f"[result] query={args.query!r} -> boxes={boxes_xyxy} scores={[round(float(s),4) for s in scores]}") |
|
|
| |
| k = min(5, num_proposals) |
| tv, ti = torch.topk(pred_scores_all.view(-1), k, dim=0) |
| print("[result] top-%d proposal scores:" % k, [round(float(v), 4) for v in tv]) |
| print("[result] top-%d proposal boxes :" % k, [ [round(x,1) for x in proposals[0][i].float().cpu().tolist()] for i in ti.cpu().tolist()]) |
|
|
| |
| draw_results( |
| image, boxes_xyxy, scores, args.query, |
| os.path.join(args.out_dir, f"pred_{stem}.png"), |
| proposals=proposals_np.tolist(), |
| proposal_path=os.path.join(args.out_dir, f"proposals_{stem}.png"), |
| ) |
| with open(os.path.join(args.out_dir, f"result_{stem}.json"), "w") as f: |
| json.dump({ |
| "image": os.path.abspath(args.image), "query": args.query, |
| "boxes_xyxy": boxes_xyxy, "scores": [float(s) for s in scores], |
| "num_proposals": num_proposals, |
| "top5_scores": [float(v) for v in tv], |
| "device": device, "dtype": str(dtype), |
| }, f, ensure_ascii=False, indent=2) |
| print("[done] all outputs saved to", os.path.abspath(args.out_dir)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
