Upload 40 files

d9c4357 verified 2 months ago

7.87 kB

#!/usr/bin/env python
"""Final verify: run multiple WeDetect-Ref-2B queries and save side-by-side panel."""
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",
]


def load_font(size):
    for p in FONT_CANDIDATES:
        if os.path.exists(p):
            try: return ImageFont.truetype(p, size)
            except: continue
    return ImageFont.load_default()


def remap_uni_keys(checkpoint):
    keys = list(checkpoint.keys())
    for key in keys:
        if "backbone" in key:
            checkpoint[key.replace("backbone.image_model.model.", "backbone.")] = checkpoint.pop(key)
    keys = list(checkpoint.keys())
    for key in keys:
        if "bbox_head" in key:
            nk = key.replace("bbox_head.head_module.", "bbox_head.")
            for a, b in [("0.2.","0.6."),("1.2.","1.6."),("2.2.","2.6."),("1.bn","4"),("1.conv","3"),("0.bn","1"),("0.conv","0")]:
                nk = nk.replace(a, b)
            checkpoint[nk] = checkpoint.pop(key)
    return checkpoint


def detect(det_model, image_path, device):
    with torch.no_grad():
        out = det_model([image_path])
    return out[0]["bboxes"].float().cpu().numpy()


def run_ref(model, processor, image, query, proposals_np, device, object_token_index, orig_w, orig_h):
    proposals = [torch.tensor(proposals_np.tolist()).to(device).to(model.dtype)]
    proposal_str = "<object>" * len(proposals[0])
    msgs = [
        {"role":"user","content":[{"type":"text","text":f'Please detect the "{query}" in the image'}]},
        {"role":"assistant","content":[{"type":"text","text":proposal_str}]},
    ]
    msgs[0]["content"] = [{"type":"image","image":copy.deepcopy(image)}] + msgs[0]["content"]
    image_inputs, video_inputs = process_vision_info(msgs, image_patch_size=16)
    texts = [processor.apply_chat_template(msgs, 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)
    t0 = time.time()
    with torch.inference_mode():
        pred = model(**model_inputs, bboxes=copy.deepcopy(proposals), ori_shapes=[(orig_w, orig_h)], bboxes_id=object_token_index)
    pos = model_inputs["input_ids"] == object_token_index
    scores = pred.logits.sigmoid()[pos].view(-1)
    return scores, time.time() - t0


def draw_panel(items, image, out_path):
    """items: list of (query, box, score)"""
    img = image.copy()
    draw = ImageDraw.Draw(img)
    font = load_font(20)
    colors = ["#FF3B30", "#FF9500", "#34C759", "#007AFF", "#AF52DE"]
    for i, (q, b, s) in enumerate(items):
        c = colors[i % len(colors)]
        x1, y1, x2, y2 = b
        draw.rectangle(((x1, y1), (x2, y2)), outline=c, width=5)
        lbl = f"Q{i+1}: {q}  score={float(s):.2f}"
        tb = draw.textbbox((0,0), lbl, font=font)
        ly = max(0, y1 - (tb[3]-tb[1]) - 8)
        draw.rectangle(((x1, ly), (x1 + (tb[2]-tb[0]) + 10, ly + (tb[3]-tb[1]) + 6)), fill=c)
        draw.text((x1+5, ly+2), lbl, fill="white", font=font)
    img.save(out_path)
    print(f"[panel] {len(items)} queries -> {out_path}")


def draw_proposals(image, props, out_path):
    img = image.copy()
    d = ImageDraw.Draw(img)
    for b in props:
        d.rectangle(((b[0], b[1]), (b[2], b[3])), outline="#00C7BE", width=2)
    img.save(out_path)
    print(f"[proposals] {len(props)} -> {out_path}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ref_dir", default="/Users/gatilin/WorkBuddy/2026-08-17-22-11-22/wedetect_verify/models/wedetect-ref-2b")
    ap.add_argument("--uni_ckpt", default="/Users/gatilin/WorkBuddy/2026-08-17-22-11-22/wedetect_verify/models/wedetect-base-uni/wedetect_base_uni.pth")
    ap.add_argument("--image", default="/Users/gatilin/WorkBuddy/2026-08-17-22-11-22/wedetect_verify/assets/demo.jpeg")
    ap.add_argument("--queries", nargs="+", default=[
        "yellow pillows on the bed",
        "the round alarm clock",
        "paintings on the wall",
        "the arc floor lamp",
        "flowers in the vase",
    ])
    ap.add_argument("--out_dir", default="/Users/gatilin/WorkBuddy/2026-08-17-22-11-22/wedetect_verify/outputs")
    ap.add_argument("--device", default="auto")
    ap.add_argument("--dtype", default="auto")
    args = ap.parse_args()

    device = args.device
    if device == "auto": device = "mps" if torch.backends.mps.is_available() else "cpu"
    dtype = torch.bfloat16 if (args.dtype == "auto" and device == "mps") else torch.float32
    print(f"[env] device={device} dtype={dtype}")
    os.makedirs(args.out_dir, exist_ok=True)

    # sanity: weights complete
    import safetensors.torch
    ref_index = json.load(open(os.path.join(args.ref_dir, "model.safetensors.index.json")))
    expected = set()
    for f, keys in ref_index.get("weight_map", {}).items():
        expected.update([f"/{k}" for k in keys])  # not used, just info
    print(f"[ref] {len(ref_index.get('weight_map', {}))} tensor entries declared across shards")

    # 1) detector
    t0 = time.time()
    det = SimpleYOLOWorldDetector(backbone_size="base", prompt_dim=768, num_prompts=256, num_proposals=100)
    ck = torch.load(args.uni_ckpt, map_location="cpu", weights_only=True)
    remap_uni_keys(ck)
    det = det.to(device).eval()
    msg = det.load_state_dict(ck, strict=False)
    print(f"[uni] loaded (missing={len(msg.missing_keys)} unexpected={len(msg.unexpected_keys)}) in {time.time()-t0:.1f}s")
    image = Image.open(args.image).convert("RGB")
    W, H = image.size
    t1 = time.time()
    props_np = detect(det, args.image, device)
    print(f"[uni] {len(props_np)} proposals in {time.time()-t1:.1f}s")
    draw_proposals(image, props_np.tolist(), os.path.join(args.out_dir, "proposals.png"))

    # 2) ref model
    t2 = time.time()
    model = Qwen3VLGroundingForConditionalGeneration.from_pretrained(args.ref_dir, torch_dtype=dtype, attn_implementation="sdpa")
    processor = AutoProcessor.from_pretrained(args.ref_dir)
    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] loaded in {time.time()-t2:.1f}s, params={sum(p.numel() for p in model.parameters())/1e9:.2f}B")

    results = []
    for q in args.queries:
        t3 = time.time()
        scores, dt = run_ref(model, processor, image, q, props_np, device, object_token_index, W, H)
        top_v, top_i = torch.topk(scores, 1, dim=0)
        box = props_np[top_i.item()].tolist()
        results.append({"query": q, "box_xyxy": box, "score": float(top_v), "forward_s": dt})
        print(f"[ref] {q!r} -> box={[round(x,1) for x in box]} score={float(top_v):.3f} ({dt:.1f}s)")

    # 3) save per-image annotated image & json
    stem = os.path.splitext(os.path.basename(args.image))[0]
    draw_panel([(r["query"], r["box_xyxy"], r["score"]) for r in results], image, os.path.join(args.out_dir, f"results_{stem}.png"))
    with open(os.path.join(args.out_dir, f"results_{stem}.json"), "w") as f:
        json.dump({"image": os.path.abspath(args.image), "device": device, "dtype": str(dtype), "results": results, "n_proposals": len(props_np)}, f, ensure_ascii=False, indent=2)
    print(f"[done] all outputs in {args.out_dir}")


if __name__ == "__main__":
    main()

	#!/usr/bin/env python
	"""Final verify: run multiple WeDetect-Ref-2B queries and save side-by-side panel."""
	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",
	]


	def load_font(size):
	for p in FONT_CANDIDATES:
	if os.path.exists(p):
	try: return ImageFont.truetype(p, size)
	except: continue
	return ImageFont.load_default()


	def remap_uni_keys(checkpoint):
	keys = list(checkpoint.keys())
	for key in keys:
	if "backbone" in key:
	checkpoint[key.replace("backbone.image_model.model.", "backbone.")] = checkpoint.pop(key)
	keys = list(checkpoint.keys())
	for key in keys:
	if "bbox_head" in key:
	nk = key.replace("bbox_head.head_module.", "bbox_head.")
	for a, b in [("0.2.","0.6."),("1.2.","1.6."),("2.2.","2.6."),("1.bn","4"),("1.conv","3"),("0.bn","1"),("0.conv","0")]:
	nk = nk.replace(a, b)
	checkpoint[nk] = checkpoint.pop(key)
	return checkpoint


	def detect(det_model, image_path, device):
	with torch.no_grad():
	out = det_model([image_path])
	return out[0]["bboxes"].float().cpu().numpy()


	def run_ref(model, processor, image, query, proposals_np, device, object_token_index, orig_w, orig_h):
	proposals = [torch.tensor(proposals_np.tolist()).to(device).to(model.dtype)]
	proposal_str = "<object>" * len(proposals[0])
	msgs = [
	{"role":"user","content":[{"type":"text","text":f'Please detect the "{query}" in the image'}]},
	{"role":"assistant","content":[{"type":"text","text":proposal_str}]},
	]
	msgs[0]["content"] = [{"type":"image","image":copy.deepcopy(image)}] + msgs[0]["content"]
	image_inputs, video_inputs = process_vision_info(msgs, image_patch_size=16)
	texts = [processor.apply_chat_template(msgs, 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)
	t0 = time.time()
	with torch.inference_mode():
	pred = model(**model_inputs, bboxes=copy.deepcopy(proposals), ori_shapes=[(orig_w, orig_h)], bboxes_id=object_token_index)
	pos = model_inputs["input_ids"] == object_token_index
	scores = pred.logits.sigmoid()[pos].view(-1)
	return scores, time.time() - t0


	def draw_panel(items, image, out_path):
	"""items: list of (query, box, score)"""
	img = image.copy()
	draw = ImageDraw.Draw(img)
	font = load_font(20)
	colors = ["#FF3B30", "#FF9500", "#34C759", "#007AFF", "#AF52DE"]
	for i, (q, b, s) in enumerate(items):
	c = colors[i % len(colors)]
	x1, y1, x2, y2 = b
	draw.rectangle(((x1, y1), (x2, y2)), outline=c, width=5)
	lbl = f"Q{i+1}: {q} score={float(s):.2f}"
	tb = draw.textbbox((0,0), lbl, font=font)
	ly = max(0, y1 - (tb[3]-tb[1]) - 8)
	draw.rectangle(((x1, ly), (x1 + (tb[2]-tb[0]) + 10, ly + (tb[3]-tb[1]) + 6)), fill=c)
	draw.text((x1+5, ly+2), lbl, fill="white", font=font)
	img.save(out_path)
	print(f"[panel] {len(items)} queries -> {out_path}")


	def draw_proposals(image, props, out_path):
	img = image.copy()
	d = ImageDraw.Draw(img)
	for b in props:
	d.rectangle(((b[0], b[1]), (b[2], b[3])), outline="#00C7BE", width=2)
	img.save(out_path)
	print(f"[proposals] {len(props)} -> {out_path}")


	def main():
	ap = argparse.ArgumentParser()
	ap.add_argument("--ref_dir", default="/Users/gatilin/WorkBuddy/2026-08-17-22-11-22/wedetect_verify/models/wedetect-ref-2b")
	ap.add_argument("--uni_ckpt", default="/Users/gatilin/WorkBuddy/2026-08-17-22-11-22/wedetect_verify/models/wedetect-base-uni/wedetect_base_uni.pth")
	ap.add_argument("--image", default="/Users/gatilin/WorkBuddy/2026-08-17-22-11-22/wedetect_verify/assets/demo.jpeg")
	ap.add_argument("--queries", nargs="+", default=[
	"yellow pillows on the bed",
	"the round alarm clock",
	"paintings on the wall",
	"the arc floor lamp",
	"flowers in the vase",
	])
	ap.add_argument("--out_dir", default="/Users/gatilin/WorkBuddy/2026-08-17-22-11-22/wedetect_verify/outputs")
	ap.add_argument("--device", default="auto")
	ap.add_argument("--dtype", default="auto")
	args = ap.parse_args()

	device = args.device
	if device == "auto": device = "mps" if torch.backends.mps.is_available() else "cpu"
	dtype = torch.bfloat16 if (args.dtype == "auto" and device == "mps") else torch.float32
	print(f"[env] device={device} dtype={dtype}")
	os.makedirs(args.out_dir, exist_ok=True)

	# sanity: weights complete
	import safetensors.torch
	ref_index = json.load(open(os.path.join(args.ref_dir, "model.safetensors.index.json")))
	expected = set()
	for f, keys in ref_index.get("weight_map", {}).items():
	expected.update([f"/{k}" for k in keys]) # not used, just info
	print(f"[ref] {len(ref_index.get('weight_map', {}))} tensor entries declared across shards")

	# 1) detector
	t0 = time.time()
	det = SimpleYOLOWorldDetector(backbone_size="base", prompt_dim=768, num_prompts=256, num_proposals=100)
	ck = torch.load(args.uni_ckpt, map_location="cpu", weights_only=True)
	remap_uni_keys(ck)
	det = det.to(device).eval()
	msg = det.load_state_dict(ck, strict=False)
	print(f"[uni] loaded (missing={len(msg.missing_keys)} unexpected={len(msg.unexpected_keys)}) in {time.time()-t0:.1f}s")
	image = Image.open(args.image).convert("RGB")
	W, H = image.size
	t1 = time.time()
	props_np = detect(det, args.image, device)
	print(f"[uni] {len(props_np)} proposals in {time.time()-t1:.1f}s")
	draw_proposals(image, props_np.tolist(), os.path.join(args.out_dir, "proposals.png"))

	# 2) ref model
	t2 = time.time()
	model = Qwen3VLGroundingForConditionalGeneration.from_pretrained(args.ref_dir, torch_dtype=dtype, attn_implementation="sdpa")
	processor = AutoProcessor.from_pretrained(args.ref_dir)
	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] loaded in {time.time()-t2:.1f}s, params={sum(p.numel() for p in model.parameters())/1e9:.2f}B")

	results = []
	for q in args.queries:
	t3 = time.time()
	scores, dt = run_ref(model, processor, image, q, props_np, device, object_token_index, W, H)
	top_v, top_i = torch.topk(scores, 1, dim=0)
	box = props_np[top_i.item()].tolist()
	results.append({"query": q, "box_xyxy": box, "score": float(top_v), "forward_s": dt})
	print(f"[ref] {q!r} -> box={[round(x,1) for x in box]} score={float(top_v):.3f} ({dt:.1f}s)")

	# 3) save per-image annotated image & json
	stem = os.path.splitext(os.path.basename(args.image))[0]
	draw_panel([(r["query"], r["box_xyxy"], r["score"]) for r in results], image, os.path.join(args.out_dir, f"results_{stem}.png"))
	with open(os.path.join(args.out_dir, f"results_{stem}.json"), "w") as f:
	json.dump({"image": os.path.abspath(args.image), "device": device, "dtype": str(dtype), "results": results, "n_proposals": len(props_np)}, f, ensure_ascii=False, indent=2)
	print(f"[done] all outputs in {args.out_dir}")


	if __name__ == "__main__":
	main()