Remove detectron2 dependency for inference

- Replace detectron2.utils.registry.Registry with minimal implementation
- Replace detectron2 transforms with pure numpy/cv2 resize+pad
- Make datasets/__init__.py lazy-import poly_data (training only)

Made-with: Cursor

app.py

@@ -11,8 +11,6 @@ from PIL import Image
 from shapely.geometry import Polygon
 
 from datasets.discrete_tokenizer import DiscreteTokenizer
-from datasets.transforms import ResizeAndPad
-from detectron2.data import transforms as T
 from models import build_model
 from util.plot_utils import plot_semantic_rich_floorplan_opencv
 
@@ -219,23 +217,23 @@ print("Loading model...")
 MODEL = load_model()
 print("Model loaded.")
 
-DATA_TRANSFORM = T.AugmentationList(
-    [ResizeAndPad((MODEL_ARGS.image_size, MODEL_ARGS.image_size), pad_value=255)]
-)
-
-
 def preprocess_image(pil_image: Image.Image) -> torch.Tensor:
+    """Resize preserving aspect ratio + pad to (image_size, image_size)."""
+    target = MODEL_ARGS.image_size
     image_np = np.array(pil_image.convert("RGB"))
-    aug_input = T.AugInput(image_np)
-    DATA_TRANSFORM(aug_input)
-    image_np = aug_input.image
+    h, w = image_np.shape[:2]
 
-    if len(image_np.shape) == 2:
-        tensor = np.expand_dims(image_np, 0)
-    else:
-        tensor = image_np.transpose((2, 0, 1))
+    scale = min(target / h, target / w)
+    new_h, new_w = int(h * scale), int(w * scale)
+    resized = cv2.resize(image_np, (new_w, new_h), interpolation=cv2.INTER_CUBIC)
+
+    padded = np.full((target, target, 3), 255, dtype=np.uint8)
+    top = (target - new_h) // 2
+    left = (target - new_w) // 2
+    padded[top:top + new_h, left:left + new_w] = resized
 
-    return (1 / 255) * torch.as_tensor(tensor, dtype=torch.float32)
+    tensor = padded.transpose((2, 0, 1)).astype(np.float32) / 255.0
+    return torch.as_tensor(tensor)
 
 
 def predict_floorplan(image: Image.Image):

datasets/__init__.py

@@ -1,7 +1,6 @@
-from .poly_data import build as build_poly
-
-
 def build_dataset(image_set, args):
+    from .poly_data import build as build_poly
+
     if args.dataset_name in ["stru3d", "cubicasa", "waffle", "r2g"]:
         print(f"Build {args.dataset_name} {image_set} dataset")
         return build_poly(image_set, args)

util/bf_utils.py

@@ -6,13 +6,26 @@ import numpy as np
 import torch
 from torch import nn
 
-from detectron2.utils.registry import Registry
-
-# need an easier place to avoid circular dependencies.
-POLY_LOSS_REGISTRY = Registry("POLY_LOSS")
-POLY_LOSS_REGISTRY.__doc__ = """
-Registry for loss computations on predicted polygons.
-"""
+class _Registry:
+    """Minimal replacement for detectron2.utils.registry.Registry."""
+    def __init__(self, name):
+        self._name = name
+        self._obj_map = {}
+
+    def register(self, obj=None):
+        if obj is None:
+            def decorator(func_or_class):
+                self._obj_map[func_or_class.__name__] = func_or_class
+                return func_or_class
+            return decorator
+        self._obj_map[obj.__name__] = obj
+        return obj
+
+    def get(self, name):
+        return self._obj_map[name]
+
+
+POLY_LOSS_REGISTRY = _Registry("POLY_LOSS")
 
 
 def box_cxcywh_to_xyxy(x):