File size: 3,974 Bytes
d9c4357 | 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 | #!/usr/bin/env python
"""Verify WeDetect-Base-Uni (proposal generator) on its own end-to-end on MPS."""
import os, sys, time, math
import torch
from PIL import Image, ImageDraw, ImageFont
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from generate_proposal import SimpleYOLOWorldDetector
CKP = "/Users/gatilin/WorkBuddy/2026-08-17-22-11-22/wedetect_verify/models/wedetect-base-uni/wedetect_base_uni.pth"
IMG = "/Users/gatilin/WorkBuddy/2026-08-17-22-11-22/wedetect_verify/assets/demo.jpeg"
OUT = "/Users/gatilin/WorkBuddy/2026-08-17-22-11-22/wedetect_verify/outputs"
def remap_uni_keys(ck):
keys = list(ck.keys())
for k in keys:
if "backbone" in k:
ck[k.replace("backbone.image_model.model.", "backbone.")] = ck.pop(k)
keys = list(ck.keys())
for k in keys:
if "bbox_head" in k:
nk = k.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)
ck[nk] = ck.pop(k)
return ck
def main():
os.makedirs(OUT, exist_ok=True)
device = "mps" if torch.backends.mps.is_available() else "cpu"
print(f"[uni] device={device}")
t0 = time.time()
det = SimpleYOLOWorldDetector(backbone_size="base", prompt_dim=768, num_prompts=256, num_proposals=100)
ck = torch.load(CKP, 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")
if msg.missing_keys[:3]: print(" missing sample:", msg.missing_keys[:3])
t1 = time.time()
with torch.no_grad():
out = det([IMG])
props = out[0]["bboxes"].float().cpu().numpy()
scores = out[0].get("scores", None)
if scores is not None: scores = scores.float().cpu().numpy()
print(f"[uni] 100 proposals in {time.time()-t1:.1f}s; shape={props.shape}; score range={scores.min() if scores is not None else 'N/A'}..{scores.max() if scores is not None else 'N/A'}")
# visualize all 100 proposals
img = Image.open(IMG).convert("RGB")
overlay = img.copy()
d = ImageDraw.Draw(overlay)
for i, (b, s) in enumerate(zip(props.tolist(), scores.tolist() if scores is not None else [None]*len(props))):
d.rectangle(((b[0], b[1]), (b[2], b[3])), outline="#00C7BE", width=2)
overlay.save(os.path.join(OUT, "uni_proposals.png"))
print(f"[vis] all 100 proposals -> {OUT}/uni_proposals.png")
# top 10 visualization
order = scores.argsort()[::-1] if scores is not None else range(len(props))
top10 = order[:10]
print("[uni] top 10 proposals:")
for i, idx in enumerate(top10):
b = props[idx].tolist()
s = scores[idx] if scores is not None else None
area = (b[2]-b[0]) * (b[3]-b[1])
print(f" #{i+1}: box=[{b[0]:.0f},{b[1]:.0f},{b[2]:.0f},{b[3]:.0f}] area={area:.0f}{' score='+(f'{s:.3f}' if s is not None else '')}")
# Save proposals.txt
out_path = os.path.join(OUT, "uni_proposals.txt")
with open(out_path, "w") as f:
f.write(f"image: {IMG}\n")
f.write(f"proposals: {len(props)} (xyxy in original image coordinates, image size = {img.size})\n")
if scores is not None: f.write("--sorted by score (desc)--\n")
if scores is not None:
for i in order:
b = props[i].tolist()
f.write(f" {i:3d} box=[{b[0]:6.1f},{b[1]:6.1f},{b[2]:6.1f},{b[3]:6.1f}] score={scores[i]:.4f} area={int((b[2]-b[0])*(b[3]-b[1]))}\n")
else:
for i, b in enumerate(props.tolist()):
f.write(f" {i:3d} box=[{b[0]:6.1f},{b[1]:6.1f},{b[2]:6.1f},{b[3]:6.1f}] area={int((b[2]-b[0])*(b[3]-b[1]))}\n")
print(f"[done] saved to {out_path}")
if __name__ == "__main__":
main()
|