ShadowNS/floorplan

1"""2Segmenter: runs CubiCasa5K if weights + the model code are available, otherwise3falls back to a pure-OpenCV heuristic so the Space is always functional.4 5CubiCasa5K background6---------------------7The published network (hg_furukawa_original) outputs a tensor whose channels8cover, in order: heatmaps, then 12 room-type segmentation channels, then 119icon channels (doors/windows/etc). We take argmax over the room channels to get10a per-pixel room-type map and treat the "Wall" class as walls. To turn the11room-type map into individual room *instances* (what the frontend needs as12labels), we run connected components on the non-wall, non-background area.13 14To actually enable the model on your Space:15  1. Add the CubiCasa5k repo's `floortrans/` package to this backend.16  2. Drop the trained weights at model/model_best_val_loss_var.pkl17     (downloadable from the CubiCasa5K project).18  3. Set USE_MODEL=1 in the Space environment.19Until then it runs the heuristic, which matches the frontend's offline engine.20"""21 22import os23 24import cv225import numpy as np26 27ROOM_CLASSES = [28    "Background", "Outdoor", "Wall", "Kitchen", "Living Room", "Bedroom",29    "Bath", "Hallway", "Railing", "Storage", "Garage", "Other",30]31WALL_ID = ROOM_CLASSES.index("Wall")32 33 34class Segmenter:35    def __init__(self, weights_path: str):36        self.weights_path = weights_path37        self.ready = False38        self.using_model = False39        self._net = None40        self._torch = None41 42    def load(self):43        want_model = os.environ.get("USE_MODEL", "0") == "1"44        if want_model and os.path.exists(self.weights_path):45            try:46                self._load_model()47                self.using_model = True48            except Exception as exc:  # noqa: BLE00149                print(f"[segmenter] model load failed, using heuristic: {exc}")50                self.using_model = False51        else:52            if want_model:53                print(f"[segmenter] weights not found at {self.weights_path}; using heuristic")54            self.using_model = False55        self.ready = True56 57    def _load_model(self):58        import torch  # local import so heuristic mode needs no torch59        from floortrans.models import get_model  # from the CubiCasa5k repo60 61        self._torch = torch62        net = get_model("hg_furukawa_original", 51)63        # CubiCasa5K final layers for 44-channel output:64        net.conv4_ = torch.nn.Conv2d(256, 44, kernel_size=1)65        net.upsample = torch.nn.ConvTranspose2d(44, 44, kernel_size=4, stride=4)66        checkpoint = torch.load(self.weights_path, map_location="cpu")67        net.load_state_dict(checkpoint["model_state"])68        net.eval()69        self._net = net70        print("[segmenter] CubiCasa5K weights loaded")71 72    # -- public API ---------------------------------------------------------73 74    def segment(self, bgr: np.ndarray):75        """Return (wall_mask uint8, room_label int32 map, n_labels)."""76        if self.using_model and self._net is not None:77            try:78                return self._segment_model(bgr)79            except Exception as exc:  # noqa: BLE00180                print(f"[segmenter] model inference failed, falling back: {exc}")81        return self._segment_heuristic(bgr)82 83    # -- model path ---------------------------------------------------------84 85    def _segment_model(self, bgr: np.ndarray):86        torch = self._torch87        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)88        h, w = rgb.shape[:2]89        # CubiCasa expects values in [0,1], shape [1,3,H,W].90        t = torch.tensor(rgb).permute(2, 0, 1).float().unsqueeze(0) / 255.091        with torch.no_grad():92            pred = self._net(t)93        if isinstance(pred, (list, tuple)):94            pred = pred[0]95        pred = torch.nn.functional.interpolate(96            pred, size=(h, w), mode="bilinear", align_corners=False97        )[0]98        # Room-type channels are the 12 after the 21 heatmap channels.99        rooms_logits = pred[21:21 + len(ROOM_CLASSES)]100        room_type = rooms_logits.argmax(0).cpu().numpy().astype(np.int32)101 102        wall_mask = ((room_type == WALL_ID).astype(np.uint8)) * 255103        labels, n = self._instances_from_roomtype(room_type)104        return wall_mask, labels, n105 106    @staticmethod107    def _instances_from_roomtype(room_type: np.ndarray):108        # Everything that is not wall and not background becomes candidate room area.109        free = ((room_type != WALL_ID) & (room_type != 0)).astype(np.uint8)110        free = cv2.morphologyEx(111            free, cv2.MORPH_OPEN,112            cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5)),113        )114        n, labels = cv2.connectedComponents(free, connectivity=8)115        return labels.astype(np.int32), n116 117    # -- heuristic path (mirrors the frontend) ------------------------------118 119    def _segment_heuristic(self, bgr: np.ndarray):120        gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)121        gray = cv2.medianBlur(gray, 3)122        _, binv = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)123 124        k = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))125        walls = cv2.morphologyEx(binv, cv2.MORPH_OPEN, k, iterations=2)126 127        free = cv2.bitwise_not(binv)128        n, labels = cv2.connectedComponents(free, connectivity=8)129        return walls, labels.astype(np.int32), n130