| |
| """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'}") |
|
|
| |
| 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") |
|
|
| |
| 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 '')}") |
|
|
| |
| 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() |
|
|
