| |
| |
| """Inference for the V3 RF-DETR Small floor-plan door-opening checkpoint. |
| Accepts a PDF (explicit --pages required) or one image. Outputs unreviewed |
| rendered-page pixel boxes, CSV, JSON, and annotated PNGs. See README.md. |
| Release adaptation: portable defaults, filename-only output provenance, |
| and the backend's 40-million-pixel PDF limit; detection functions unchanged. |
| """ |
| import argparse |
| import csv |
| from datetime import datetime |
| from importlib.metadata import PackageNotFoundError, version |
| import json |
| import math |
| from pathlib import Path |
|
|
| import numpy as np |
| from PIL import Image, ImageDraw, ImageFont, ImageOps |
|
|
| TILE, STRIDE = 768, 384 |
|
|
|
|
| def positions(length): |
| return sorted(set(range(0, max(1, length - TILE + 1), STRIDE)) | {max(0, length - TILE)}) |
|
|
|
|
| def ious(box, boxes): |
| if not len(boxes): |
| return np.empty(0) |
| lo = np.maximum(box[:2], boxes[:, :2]) |
| hi = np.minimum(box[2:], boxes[:, 2:]) |
| inter = np.maximum(hi - lo, 0).prod(axis=1) |
| area = np.maximum(box[2:] - box[:2], 0).prod() |
| areas = np.maximum(boxes[:, 2:] - boxes[:, :2], 0).prod(axis=1) |
| return inter / np.maximum(area + areas - inter, 1e-9) |
|
|
|
|
| def nms(boxes, scores, threshold): |
| order = np.argsort(-scores, kind='stable') |
| keep = [] |
| while len(order): |
| idx = int(order[0]) |
| keep.append(idx) |
| rest = order[1:] |
| order = rest[ious(boxes[idx], boxes[rest]) <= threshold] |
| return np.asarray(keep, dtype=int) |
|
|
|
|
| def edge_sides(box, origin, size, margin=4): |
| x, y = origin |
| w, h = size |
| return [side for side, flag in enumerate(( |
| x > 0 and box[0] - x <= margin, |
| y > 0 and box[1] - y <= margin, |
| x + TILE < w and x + TILE - box[2] <= margin, |
| y + TILE < h and y + TILE - box[3] <= margin)) if flag] |
|
|
|
|
| def confirms(fragment, complete): |
| if not fragment['edges'] or complete['edges']: |
| return False |
| a = np.array(fragment['box']) |
| b = np.array(complete['box']) |
| area = np.prod(a[2:] - a[:2]) |
| inter = np.maximum(np.minimum(a[2:], b[2:]) - np.maximum(a[:2], b[:2]), 0).prod() |
| if area <= 0 or inter / area < .85: |
| return False |
| x, y = fragment['origin'] |
| bounds = np.array([x, y, x + TILE, y + TILE]) |
| if not all((b[k] < bounds[k] - 4 if k < 2 else b[k] > bounds[k] + 4) for k in fragment['edges']): |
| return False |
| clipped = np.r_[np.maximum(b[:2], bounds[:2]), np.minimum(b[2:], bounds[2:])] |
| return bool(ious(a, clipped.reshape(1, 4))[0] >= .5) |
|
|
|
|
| def detect_tile(model, page, origin, threshold): |
| x, y = origin |
| w, h = page.size |
| tile = Image.new('RGB', (TILE, TILE), 'white') |
| tile.paste(page.crop((x, y, min(x + TILE, w), min(y + TILE, h))), (0, 0)) |
| detections = model.predict(tile, threshold=threshold) |
| result = [] |
| for box, score in zip(detections.xyxy, detections.confidence): |
| box = np.array(box, dtype=float) + np.array([x, y, x, y]) |
| box[[0, 2]] = np.clip(box[[0, 2]], 0, w) |
| box[[1, 3]] = np.clip(box[[1, 3]], 0, h) |
| if np.isfinite(box).all() and np.isfinite(score) and np.all(box[2:] > box[:2]): |
| result.append(dict(box=box.tolist(), score=float(score), origin=list(origin), |
| edges=edge_sides(box, origin, page.size))) |
| return result |
|
|
|
|
| def detect_page(model, page, threshold, nms_iou): |
| w, h = page.size |
| grid = [(x, y) for y in positions(h) for x in positions(w)] |
| cache, raw = {}, [] |
| print(f' {len(grid)} initial tiles', flush=True) |
| for number, origin in enumerate(grid, 1): |
| cache[origin] = detect_tile(model, page, origin, threshold) |
| raw.extend(cache[origin]) |
| if number % 20 == 0 or number == len(grid): |
| print(f' Tile {number}/{len(grid)}', flush=True) |
| fragments = [r for r in raw if r['edges']] |
| complete = [r for r in raw if not r['edges']] |
| pending = [r for r in fragments if not any(confirms(r, c) for c in complete)] |
| origins = set() |
| for r in pending: |
| b = r['box'] |
| cx, cy = (b[0] + b[2]) / 2, (b[1] + b[3]) / 2 |
| origins.add((int(np.clip(round((cx - TILE / 2) / 32) * 32, 0, max(0, w - TILE))), |
| int(np.clip(round((cy - TILE / 2) / 32) * 32, 0, max(0, h - TILE))))) |
| origins = sorted(origins - set(cache)) |
| print(f' Rechecking {len(origins)} shifted crops', flush=True) |
| added = [] |
| for number, origin in enumerate(origins, 1): |
| ds = detect_tile(model, page, origin, threshold) |
| added.extend(c for c in ds if not c['edges'] and any(confirms(r, c) for r in pending)) |
| if number % 20 == 0 or number == len(origins): |
| print(f' Recheck {number}/{len(origins)}', flush=True) |
| candidates = complete + added |
| removed = {i for i, r in enumerate(raw) if r['edges'] and any(confirms(r, c) for c in candidates)} |
| revised = [r for i, r in enumerate(raw) if i not in removed] + added |
| boxes = np.array([r['box'] for r in revised], float).reshape(-1, 4) |
| scores = np.array([r['score'] for r in revised], float) |
| keep = nms(boxes, scores, nms_iou) |
| return [revised[i] for i in keep], dict(initial_tiles=len(grid), shifted_crops=len(origins), |
| suppressed_fragments=len(removed)) |
|
|
|
|
| def page_numbers(text, count): |
| selected = set() |
| for group in text.split(','): |
| bounds = [int(v) for v in group.strip().split('-')] |
| if len(bounds) == 1: |
| selected.add(bounds[0]) |
| elif len(bounds) == 2 and bounds[0] <= bounds[1]: |
| selected.update(range(bounds[0], bounds[1] + 1)) |
| else: |
| raise ValueError('Use PDF page numbers such as 5 or 5,7-9.') |
| if not selected or min(selected) < 1 or max(selected) > count: |
| raise ValueError(f'PDF page numbers must be between 1 and {count}.') |
| return sorted(selected) |
|
|
|
|
| def roi_pixels(size, roi): |
| w, h = size |
| if roi is None: |
| return (0, 0, w, h) |
| left, top, right, bottom = roi |
| return (math.floor(left * w / 100), math.floor(top * h / 100), |
| math.ceil(right * w / 100), math.ceil(bottom * h / 100)) |
|
|
|
|
| def label_font(size): |
| for name in ('Arial.ttf', 'DejaVuSans.ttf'): |
| try: |
| return ImageFont.truetype(name, size) |
| except OSError: |
| pass |
| return ImageFont.load_default() |
|
|
|
|
| def annotated_image(page, detections, roi, preview=False): |
| image = page.copy() |
| draw = ImageDraw.Draw(image) |
| scale = max(.5, min(2, page.width / 6300)) |
| width = max(2, round(4 * scale)) |
| font = label_font(max(14, round(23 * scale))) |
| if roi != (0, 0, page.width, page.height): |
| draw.rectangle((roi[0], roi[1], roi[2] - 1, roi[3] - 1), outline='#dd8700', width=width) |
| for d in detections: |
| box = [round(v) for v in d['box']] |
| draw.rectangle(box, outline='#0066ff', width=width) |
| position = (max(0, box[0]), max(0, box[1] - round(28 * scale))) |
| label = f'{d["id"]} {d["score"]:.2f}' |
| draw.rectangle(draw.textbbox(position, label, font=font), fill='white') |
| draw.text(position, label, fill='#0066ff', font=font) |
| title = 'PREVIEW ONLY - no model run' if preview else f'Detections: {len(detections)} - unreviewed' |
| draw.rectangle(draw.textbbox((10, 10), title, font=font), fill='white') |
| draw.text((10, 10), title, font=font, fill='#0066ff') |
| return image |
|
|
|
|
| def save_report(output, report): |
| path = output / 'predictions.json' |
| temporary = output / 'predictions.json.tmp' |
| temporary.write_text(json.dumps(report, indent=2) + '\n') |
| temporary.replace(path) |
| with (output / 'detections.csv').open('w', newline='') as f: |
| writer = csv.writer(f) |
| writer.writerow(['input', 'pdf_page', 'detection_id', 'confidence', 'x1', 'y1', 'x2', 'y2']) |
| for p in report['pages']: |
| for d in p['detections']: |
| writer.writerow([report['input'], p['pdf_page'], d['id'], d['score'], *d['box']]) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| parser.add_argument('--input', required=True, type=Path, help='PDF or one PNG/JPEG image.') |
| parser.add_argument('--pages', help='Required for PDFs: 1-based page numbers, e.g. 5 or 5,7-9.') |
| parser.add_argument('--checkpoint', type=Path, default=Path(__file__).resolve().parent / 'model.safetensors') |
| parser.add_argument('--model-config', type=Path, help='Architecture JSON; defaults beside weights.') |
| parser.add_argument('--output', type=Path, help='New output folder; defaults to a timestamped outputs/ folder.') |
| parser.add_argument('--dpi', type=int, default=150, help='PDF rendering DPI. Default: 150, matching training preparation.') |
| parser.add_argument('--threshold', type=float, default=.3) |
| parser.add_argument('--nms-iou', type=float, default=.5) |
| parser.add_argument('--device', choices=['mps', 'cuda', 'cpu'], default='cpu') |
| parser.add_argument('--roi', nargs=4, type=float, metavar=('LEFT', 'TOP', 'RIGHT', 'BOTTOM'), |
| help='Drawing-area bounds as percentages, each 0..100. Same bounds apply to every selected page.') |
| parser.add_argument('--preview-only', action='store_true', help='Render selected pages and ROI without loading the model.') |
| args = parser.parse_args() |
| if not args.input.is_file(): |
| parser.error(f'Input missing: {args.input}') |
| for key in ('threshold', 'nms_iou'): |
| if not 0 < getattr(args, key) <= 1: |
| parser.error(f'{key} must be in (0, 1].') |
| if not 36 <= args.dpi <= 600: |
| parser.error('--dpi must be between 36 and 600.') |
| if args.roi and not (all(math.isfinite(v) for v in args.roi) and |
| 0 <= args.roi[0] < args.roi[2] <= 100 and 0 <= args.roi[1] < args.roi[3] <= 100): |
| parser.error('--roi needs LEFT < RIGHT and TOP < BOTTOM, within 0..100.') |
| output = args.output or Path('outputs') / ('door_predictions_' + datetime.now().strftime('%Y%m%d_%H%M%S_%f')) |
| if output.exists(): |
| parser.error('Output already exists. Use a new --output folder.') |
| document, source_image = None, None |
| try: |
| if args.input.suffix.lower() == '.pdf': |
| if not args.pages: |
| parser.error('Choose PDF pages with --pages, using 1-based PDF page numbers.') |
| try: |
| import pymupdf |
| except ImportError: |
| parser.error('Install PDF support: .venv/bin/python -m pip install pymupdf') |
| document = pymupdf.open(args.input) |
| if document.needs_pass: |
| parser.error('The PDF is password-protected. Provide an unlocked copy.') |
| selected = page_numbers(args.pages, len(document)) |
| else: |
| if args.pages: |
| parser.error('--pages applies only to PDFs.') |
| with Image.open(args.input) as opened: |
| source = ImageOps.exif_transpose(opened).convert('RGBA') |
| background = Image.new('RGBA', source.size, 'white') |
| source_image = Image.alpha_composite(background, source).convert('RGB') |
| selected = [None] |
| except (OSError, ValueError) as error: |
| if document is not None: |
| document.close() |
| parser.error(str(error)) |
| model, torch = None, None |
| if not args.preview_only: |
| if not args.checkpoint.is_file(): |
| parser.error(f'Checkpoint missing: {args.checkpoint}') |
| try: |
| installed = version('rfdetr') |
| except PackageNotFoundError: |
| parser.error('Install the pinned environment: python -m pip install "rfdetr==1.11.2"') |
| if installed != '1.11.2': |
| parser.error(f'This script requires RF-DETR 1.11.2; found {installed}.') |
| import torch |
| from model_loader import load_model |
| if args.device == 'mps' and not torch.backends.mps.is_available(): |
| parser.error('MPS unavailable; use the Mac environment or pass --device cpu.') |
| if args.device == 'cuda' and not torch.cuda.is_available(): |
| parser.error('CUDA unavailable.') |
| print('Loading safetensors door weights...', flush=True) |
| model = load_model(args.checkpoint, args.model_config, device=args.device) |
| output.mkdir(parents=True) |
| report = dict(input=args.input.name, checkpoint=None if args.preview_only else args.checkpoint.name, |
| status='running', preview_only=args.preview_only, threshold=args.threshold, nms_iou=args.nms_iou, |
| tile_size=TILE, stride=STRIDE, dpi=args.dpi if document is not None else None, |
| requested_pages=selected, roi_percent=args.roi, |
| coordinate_system='Full rendered-page pixels: x rightward, y downward, origin at top left.', |
| note='Unreviewed model detections. Counts may include false detections, duplicates, and missed doors. ' |
| 'Sums across pages do not deduplicate repeated floor plans. The ROI is user-selected, not automatic.', |
| total_detections_across_selected_pages=None if args.preview_only else 0, pages=[]) |
| save_report(output, report) |
| try: |
| for number in selected: |
| stem = args.input.stem if number is None else f'{args.input.stem}_page_{number:03d}' |
| print(f'\n{stem}', flush=True) |
| if document is not None: |
| pdf_page = document[number - 1] |
| estimate = math.ceil(pdf_page.rect.width * args.dpi / 72) * math.ceil(pdf_page.rect.height * args.dpi / 72) |
| if estimate > 40_000_000: |
| raise ValueError('PDF page exceeds 40 million rendered pixels; crop a copy or lower --dpi.') |
| pix = pdf_page.get_pixmap(dpi=args.dpi, colorspace=pymupdf.csRGB, alpha=False) |
| page = Image.frombytes('RGB', (pix.width, pix.height), pix.samples) |
| else: |
| page = source_image |
| rect = roi_pixels(page.size, args.roi) |
| records, details = [], {} |
| if not args.preview_only: |
| with torch.inference_mode(): |
| records, details = detect_page(model, page.crop(rect), args.threshold, args.nms_iou) |
| detections = [] |
| for index, record in enumerate(records, 1): |
| x1, y1, x2, y2 = record['box'] |
| detections.append(dict(id=index, score=record['score'], |
| box=[x1 + rect[0], y1 + rect[1], x2 + rect[0], y2 + rect[1]])) |
| image_name = stem + ('_preview.png' if args.preview_only else '_doors.png') |
| annotated_image(page, detections, rect, args.preview_only).save(output / image_name, compress_level=1) |
| entry = dict(pdf_page=number, image_size=list(page.size), roi_pixels=list(rect), |
| count=None if args.preview_only else len(detections), detections=detections, |
| annotated_image=image_name, processing=details) |
| report['pages'].append(entry) |
| if not args.preview_only: |
| report['total_detections_across_selected_pages'] += len(detections) |
| save_report(output, report) |
| print(' Preview saved.' if args.preview_only else f' Saved {len(detections)} unreviewed door detections.', flush=True) |
| report['status'] = 'completed' |
| except KeyboardInterrupt: |
| report['status'] = 'interrupted' |
| print('\nStopped. Completed pages remain saved; the interrupted page is not complete.', flush=True) |
| except Exception as error: |
| report['status'] = 'error' |
| report['error'] = str(error) |
| raise |
| finally: |
| save_report(output, report) |
| if document is not None: |
| document.close() |
| print(f'\nResults: {output.resolve()}', flush=True) |
| if report['status'] == 'interrupted': |
| raise SystemExit(130) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
