| """WeDetect-Ref inference adapted for CPU/MPS (Apple Silicon). |
| |
| Two-stage pipeline: |
| 1. WeDetect-Uni (YOLO-World style) generates up to 100 object proposals. |
| 2. WeDetect-Ref (Qwen3-VL based) classifies proposals in a single forward pass. |
| |
| Usage: |
| python infer_ref_mac.py --ref_ckpt <path-to-WeDetect-Ref-4B> \ |
| --uni_ckpt <path-to-wedetect_base_uni.pth> \ |
| --image assets/demo.jpeg --query "the bed" --device cpu --out pred.png |
| """ |
| import argparse |
| import copy |
| import time |
|
|
| import torch |
| from PIL import Image |
| from transformers import AutoProcessor |
|
|
| from wedetect_ref.models.qwen3vl_referring import Qwen3VLGroundingForConditionalGeneration |
| from wedetect_ref.models.vision_process import process_vision_info |
| from generate_proposal import SimpleYOLOWorldDetector |
| from vis import plot_bounding_boxes |
|
|
|
|
| def remap_uni_checkpoint(checkpoint): |
| """Same key remapping as the official infer_wedetect_ref.py.""" |
| 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('--ref_ckpt', type=str, required=True) |
| parser.add_argument('--uni_ckpt', type=str, required=True) |
| parser.add_argument('--image', type=str, required=True) |
| parser.add_argument('--query', type=str, required=True) |
| parser.add_argument('--device', type=str, default='cpu', choices=['cpu', 'mps']) |
| parser.add_argument('--score_thre', type=float, default=-1.0) |
| parser.add_argument('--out', type=str, default='pred.png') |
| args = parser.parse_args() |
|
|
| device = torch.device(args.device) |
| torch.set_grad_enabled(False) |
|
|
| |
| t0 = time.time() |
| model_size = 'base' if 'base' in args.uni_ckpt else 'large' |
| det_model = SimpleYOLOWorldDetector( |
| backbone_size=model_size, prompt_dim=768, num_prompts=256, num_proposals=100) |
| checkpoint = torch.load(args.uni_ckpt, map_location='cpu') |
| checkpoint = remap_uni_checkpoint(checkpoint) |
| msg = det_model.load_state_dict(checkpoint, strict=False) |
| print(f'[uni] load_state_dict: {msg}', flush=True) |
| det_model = det_model.float().to(device).eval() |
|
|
| outputs = det_model([args.image]) |
| boxes = outputs[0]['bboxes'].float().cpu() |
| print(f'[uni] {len(boxes)} proposals in {time.time() - t0:.1f}s', flush=True) |
|
|
| |
| t0 = time.time() |
| model = Qwen3VLGroundingForConditionalGeneration.from_pretrained( |
| args.ref_ckpt, |
| torch_dtype=torch.float16, |
| attn_implementation="sdpa", |
| ) |
| processor = AutoProcessor.from_pretrained(args.ref_ckpt) |
| object_token_index = processor.tokenizer.convert_tokens_to_ids("<object>") |
| model.model.object_token_id = object_token_index |
| model = model.to(device).eval() |
| |
| for name in ['first_scale_conv1', 'first_scale_norm', 'first_scale_conv2', |
| 'second_scale_conv', 'merge', 'image_pos_projector', |
| 'object_vision_projector', 'object_pos_projector']: |
| getattr(model.model, name).to('cpu', torch.float32) |
| print(f'[ref] model loaded on {args.device} in {time.time() - t0:.1f}s', flush=True) |
|
|
| image = Image.open(args.image).convert('RGB') |
| ori_shape = [image.size] |
| proposals = [boxes.to(device).to(model.dtype)] |
|
|
| num_proposals = len(boxes) |
| messages = [ |
| { |
| "role": "user", |
| "content": [ |
| {"type": "image", "image": copy.deepcopy(image)}, |
| {"type": "text", "text": 'Please detect the "%s" in the image' % args.query}, |
| ], |
| }, |
| { |
| "role": "assistant", |
| "content": [{"type": "text", "text": "<object>" * num_proposals}], |
| }, |
| ] |
| 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(device) |
|
|
| t0 = time.time() |
| with torch.inference_mode(): |
| pred = model( |
| **model_inputs, |
| bboxes=copy.deepcopy(proposals), |
| ori_shapes=ori_shape, |
| bboxes_id=object_token_index, |
| image_inputs=image_inputs, |
| ) |
| print(f'[ref] forward pass in {time.time() - t0:.1f}s', flush=True) |
|
|
| proposal_positions = model_inputs['input_ids'] == object_token_index |
| pred_scores = pred.logits.sigmoid()[proposal_positions].view(-1).float().cpu() |
| pred_bboxes = boxes.clone() |
|
|
| if args.score_thre < 0: |
| topk_values, topk_indexes = torch.topk(pred_scores.view(-1), 1, dim=0) |
| pred_bboxes = pred_bboxes[topk_indexes] |
| pred_scores = topk_values |
| else: |
| mask = pred_scores > args.score_thre |
| pred_bboxes = pred_bboxes[mask] |
| pred_scores = pred_scores[mask] |
|
|
| print(f'[result] query="{args.query}" -> {len(pred_bboxes)} box(es)') |
| for b, s in zip(pred_bboxes.tolist(), pred_scores.tolist()): |
| print(f' score={s:.4f} box=[{b[0]:.1f}, {b[1]:.1f}, {b[2]:.1f}, {b[3]:.1f}]') |
|
|
| pred_image = plot_bounding_boxes(copy.deepcopy(image), pred_bboxes.tolist()) |
| pred_image.save(args.out) |
| print(f'[vis] saved to {args.out}') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
