enhance: preload depth model and compress output pixels payload

app.py

@@ -4,6 +4,7 @@ import io
 import json
 import os
 import shutil
+import time
 try:
     import tomllib
 except ImportError:
@@ -155,6 +156,7 @@ def _load_segmentation_model():
     if SEGMENTATION_MODEL == "oneformer":
         try:
             print(f"Loading OneFormer: {ONEFORMER_MODEL_NAME} ...", flush=True)
+            start_time = time.perf_counter()
             seg_processor = OneFormerProcessor.from_pretrained(
                 ONEFORMER_MODEL_NAME,
                 local_files_only=hf_offline(),
@@ -165,7 +167,7 @@ def _load_segmentation_model():
             ).to(device)
             seg_model.eval()
             segmentation_backend = "oneformer"
-            print("OneFormer loaded.", flush=True)
+            print(f"OneFormer loaded in {time.perf_counter() - start_time:.4f}s.", flush=True)
             return
         except Exception as exc:
             print(f"OneFormer failed ({exc}), falling back to Mask2Former.", flush=True)
@@ -173,6 +175,7 @@ def _load_segmentation_model():
     if SEGMENTATION_MODEL in {"oneformer", "mask2former"}:
         try:
             print(f"Loading Mask2Former: {MASK2FORMER_MODEL_NAME} ...", flush=True)
+            start_time = time.perf_counter()
             seg_processor = AutoImageProcessor.from_pretrained(
                 MASK2FORMER_MODEL_NAME,
                 local_files_only=hf_offline(),
@@ -183,12 +186,13 @@ def _load_segmentation_model():
             ).to(device)
             seg_model.eval()
             segmentation_backend = "mask2former"
-            print("Mask2Former loaded.", flush=True)
+            print(f"Mask2Former loaded in {time.perf_counter() - start_time:.4f}s.", flush=True)
             return
         except Exception as exc:
             print(f"Mask2Former failed ({exc}), falling back to SegFormer.", flush=True)
 
     print(f"Loading SegFormer: {SEGFORMER_MODEL_NAME} ...", flush=True)
+    start_time = time.perf_counter()
     seg_processor = AutoImageProcessor.from_pretrained(
         SEGFORMER_MODEL_NAME,
         local_files_only=hf_offline(),
@@ -199,7 +203,7 @@ def _load_segmentation_model():
     ).to(device)
     seg_model.eval()
     segmentation_backend = "segformer"
-    print("SegFormer loaded.", flush=True)
+    print(f"SegFormer loaded in {time.perf_counter() - start_time:.4f}s.", flush=True)
 
 
 _load_segmentation_model()
@@ -210,15 +214,40 @@ def _load_intrinsic_model():
     if ENABLE_INTRINSIC_SHADING and intrinsic_models is None:
         try:
             print(f"Loading Intrinsic Image Decomposition model: {INTRINSIC_MODEL_VERSION} ...", flush=True)
+            start_time = time.perf_counter()
             from intrinsic.pipeline import load_models
             intrinsic_models = load_models(INTRINSIC_MODEL_VERSION, device=str(device))
-            print("Intrinsic model loaded.", flush=True)
+            print(f"Intrinsic model loaded in {time.perf_counter() - start_time:.4f}s.", flush=True)
         except Exception as exc:
             print(f"Intrinsic model failed to load ({exc}). Falling back to luminance shading.", flush=True)
 
 
 _load_intrinsic_model()
 
+
+def _load_depth_model():
+    global depth_processor, depth_model
+    if ENABLE_DEPTH_ESTIMATION and (depth_processor is None or depth_model is None):
+        try:
+            model_name = DEPTH_MODEL_NAME
+            print(f"Loading depth model: {model_name} ...", flush=True)
+            start_time = time.perf_counter()
+            depth_processor = AutoImageProcessor.from_pretrained(
+                model_name,
+                local_files_only=hf_offline(),
+            )
+            depth_model = AutoModelForDepthEstimation.from_pretrained(
+                model_name,
+                local_files_only=hf_offline(),
+            ).to(device)
+            depth_model.eval()
+            print(f"Depth model loaded in {time.perf_counter() - start_time:.4f}s.", flush=True)
+        except Exception as exc:
+            print(f"Depth model failed to load ({exc}).", flush=True)
+
+
+_load_depth_model()
+
 app = FastAPI()
 app.add_middleware(
     CORSMiddleware,
@@ -262,15 +291,18 @@ def class_ids(names: set[str]) -> list[int]:
     return [idx for idx, name in enumerate(ADE20K_CLASSES) if name in names]
 
 
-def estimate_depth(img: Image.Image, width: int, height: int):
+def estimate_depth(img: Image.Image, width: int, height: int, task_id: str = "segment"):
     global depth_processor, depth_model
     if not ENABLE_DEPTH_ESTIMATION:
         return None
 
     model_name = DEPTH_MODEL_NAME
+    print(f"[{task_id}] Starting depth estimation...", flush=True)
+    start_time = time.perf_counter()
     try:
         if depth_processor is None or depth_model is None:
-            print(f"Loading depth model: {model_name} ...", flush=True)
+            print(f"[{task_id}] Loading depth model: {model_name} ...", flush=True)
+            start_load = time.perf_counter()
             depth_processor = AutoImageProcessor.from_pretrained(
                 model_name,
                 local_files_only=hf_offline(),
@@ -280,7 +312,7 @@ def estimate_depth(img: Image.Image, width: int, height: int):
                 local_files_only=hf_offline(),
             ).to(device)
             depth_model.eval()
-            print("Depth model loaded.", flush=True)
+            print(f"[{task_id}] Depth model loaded in {time.perf_counter() - start_load:.4f}s.", flush=True)
 
         inputs = depth_processor(images=img, return_tensors="pt").to(device)
         with torch.no_grad():
@@ -293,18 +325,22 @@ def estimate_depth(img: Image.Image, width: int, height: int):
             ).squeeze().cpu().numpy()
         depth = cv2.GaussianBlur(depth.astype(np.float32), (0, 0), sigmaX=3)
         depth_min, depth_max = float(np.min(depth)), float(np.max(depth))
+        duration = time.perf_counter() - start_time
+        print(f"[{task_id}] Depth estimation completed in {duration:.4f}s", flush=True)
         if depth_max - depth_min < 1e-6:
             return None
         return (depth - depth_min) / (depth_max - depth_min)
     except Exception as exc:
-        print(f"Depth estimation skipped ({exc}).", flush=True)
+        print(f"[{task_id}] Depth estimation skipped ({exc}).", flush=True)
         return None
 
 
-def build_shade_map(img_np: np.ndarray, surface_mask: np.ndarray) -> np.ndarray | None:
+def build_shade_map(img_np: np.ndarray, surface_mask: np.ndarray, task_id: str = "segment") -> np.ndarray | None:
     if not surface_mask.any():
         return None
 
+    print(f"[{task_id}] Starting shade map build...", flush=True)
+    start_time = time.perf_counter()
     mask = surface_mask.astype(np.uint8)
     luminance = (
         img_np[:, :, 0].astype(np.float32) * 0.299
@@ -337,13 +373,18 @@ def build_shade_map(img_np: np.ndarray, surface_mask: np.ndarray) -> np.ndarray
     smooth = cv2.GaussianBlur(filled, (0, 0), sigmaX=sigma, sigmaY=sigma)
     shade = np.clip(smooth / median_lum, 0.55, 1.35)
     shade[mask == 0] = 1.0
-    return np.round((shade - 0.55) * (255.0 / 0.80)).clip(0, 255).astype(np.uint8)
+    result = np.round((shade - 0.55) * (255.0 / 0.80)).clip(0, 255).astype(np.uint8)
+    duration = time.perf_counter() - start_time
+    print(f"[{task_id}] Shade map built in {duration:.4f}s", flush=True)
+    return result
 
 
-def build_intrinsic_shade_map(img_np: np.ndarray, surface_mask: np.ndarray) -> np.ndarray | None:
+def build_intrinsic_shade_map(img_np: np.ndarray, surface_mask: np.ndarray, task_id: str = "segment") -> np.ndarray | None:
     if not surface_mask.any() or intrinsic_models is None:
         return None
 
+    print(f"[{task_id}] Starting intrinsic shade map build...", flush=True)
+    start_time = time.perf_counter()
     try:
         # Convert image to float32 range [0.0, 1.0] as expected by compphoto/Intrinsic
         img_float = img_np.astype(np.float32) / 255.0
@@ -352,7 +393,7 @@ def build_intrinsic_shade_map(img_np: np.ndarray, surface_mask: np.ndarray) -> n
         from intrinsic.pipeline import run_pipeline
 
         # Use CPU/CUDA device string
-        results = run_pipeline(intrinsic_models, img_float, device=str(device))
+        results = run_pipeline(intrinsic_models, img_float, stage=1, device=str(device))
 
         # Extract shading map
         shading = None
@@ -402,9 +443,12 @@ def build_intrinsic_shade_map(img_np: np.ndarray, surface_mask: np.ndarray) -> n
         relative_shading[surface_mask == 0] = 1.0
 
         # Encode to [0, 255] byte range matching the frontend
-        return np.round((relative_shading - 0.55) * (255.0 / 0.80)).clip(0, 255).astype(np.uint8)
+        result = np.round((relative_shading - 0.55) * (255.0 / 0.80)).clip(0, 255).astype(np.uint8)
+        duration = time.perf_counter() - start_time
+        print(f"[{task_id}] Intrinsic shade map built in {duration:.4f}s", flush=True)
+        return result
     except Exception as exc:
-        print(f"Intrinsic shading decomposition failed: {exc}. Falling back to default luminance shading.", flush=True)
+        print(f"[{task_id}] Intrinsic shading decomposition failed: {exc}. Falling back to default luminance shading.", flush=True)
         return None
 
 
@@ -532,7 +576,9 @@ def detect_vanishing_point(img_np: np.ndarray, floor_mask: np.ndarray):
     return {"x": float(center[0]), "y": float(center[1])}
 
 
-def estimate_floor_plane(mask: np.ndarray, img_np: np.ndarray):
+def estimate_floor_plane(mask: np.ndarray, img_np: np.ndarray, task_id: str = "segment"):
+    print(f"[{task_id}] Starting floor plane estimation...", flush=True)
+    start_time = time.perf_counter()
     ys, xs = np.where(mask > 0)
     if len(xs) < 1000:
         return None, None
@@ -593,6 +639,8 @@ def estimate_floor_plane(mask: np.ndarray, img_np: np.ndarray):
         return None, None
     dst = np.float32([[x1, y2], [x2, y2], [x2, y1], [x1, y1]])
     homography = cv2.getPerspectiveTransform(src, dst).flatten().tolist()
+    duration = time.perf_counter() - start_time
+    print(f"[{task_id}] Floor plane estimation completed in {duration:.4f}s", flush=True)
     return homography, {
         "x": x1,
         "y": y1,
@@ -608,7 +656,10 @@ def build_floor_surface_mask(
     seg_map: np.ndarray,
     quad: np.ndarray | None,
     depth: np.ndarray | None,
+    task_id: str = "segment",
 ):
+    print(f"[{task_id}] Starting floor surface mask build...", flush=True)
+    start_time = time.perf_counter()
     h, w = floor_mask.shape[:2]
     kern_size = max(5, min(h, w) // 160) | 1
     kern = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kern_size, kern_size))
@@ -655,11 +706,15 @@ def build_floor_surface_mask(
     surface = cv2.dilate(surface, boundary_kern, iterations=1)
     surface[occ_dilated > 0] = 0
     surface[reject_dilated > 0] = 0
+    duration = time.perf_counter() - start_time
+    print(f"[{task_id}] Floor surface mask built in {duration:.4f}s", flush=True)
     return surface
 
 
-def run_segmentation(img: Image.Image, img_np: np.ndarray):
+def run_segmentation(img: Image.Image, img_np: np.ndarray, task_id: str = "segment"):
     h, w = img_np.shape[:2]
+    print(f"[{task_id}] Running segmentation (backend: {segmentation_backend})...", flush=True)
+    start_time = time.perf_counter()
     if segmentation_backend == "oneformer":
         inputs = seg_processor(
             images=img,
@@ -672,9 +727,9 @@ def run_segmentation(img: Image.Image, img_np: np.ndarray):
             outputs,
             target_sizes=[(h, w)],
         )[0]
-        return result.cpu().numpy().astype(np.uint8)
+        seg_map = result.cpu().numpy().astype(np.uint8)
 
-    if segmentation_backend == "mask2former":
+    elif segmentation_backend == "mask2former":
         inputs = seg_processor(images=img, return_tensors="pt").to(device)
         with torch.no_grad():
             outputs = seg_model(**inputs)
@@ -688,18 +743,23 @@ def run_segmentation(img: Image.Image, img_np: np.ndarray):
             pan_map = pan_result["segmentation"].cpu().numpy()
             for seg_info in pan_result["segments_info"]:
                 seg_map[pan_map == seg_info["id"]] = min(seg_info["label_id"], 255)
-            return seg_map
-        result = seg_processor.post_process_semantic_segmentation(
-            outputs,
-            target_sizes=[(h, w)],
-        )[0]
-        return result.cpu().numpy().astype(np.uint8)
+        else:
+            result = seg_processor.post_process_semantic_segmentation(
+                outputs,
+                target_sizes=[(h, w)],
+            )[0]
+            seg_map = result.cpu().numpy().astype(np.uint8)
+
+    else:
+        inputs = seg_processor(images=img, return_tensors="pt").to(device)
+        with torch.no_grad():
+            outputs = seg_model(**inputs)
+        seg = outputs.logits.argmax(dim=1).squeeze().cpu().numpy()
+        seg_map = cv2.resize(seg.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST)
 
-    inputs = seg_processor(images=img, return_tensors="pt").to(device)
-    with torch.no_grad():
-        outputs = seg_model(**inputs)
-    seg = outputs.logits.argmax(dim=1).squeeze().cpu().numpy()
-    return cv2.resize(seg.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST)
+    duration = time.perf_counter() - start_time
+    print(f"[{task_id}] Segmentation completed in {duration:.4f}s", flush=True)
+    return seg_map
 
 
 def segmenter_metadata_name() -> str:
@@ -708,15 +768,16 @@ def segmenter_metadata_name() -> str:
     return segmentation_backend
 
 
-def build_segmentation_bundle(contents: bytes):
+def build_segmentation_bundle(contents: bytes, task_id: str = "segment"):
+    print(f"[{task_id}] Starting bundle build...", flush=True)
+    start_total = time.perf_counter()
     img = Image.open(io.BytesIO(contents)).convert("RGB")
     img_np = np.array(img)
     h, w = img_np.shape[:2]
     min_floor_area = max(1200, int(w * h * 0.015))
 
-    seg_map = run_segmentation(img, img_np)
-    rgba = np.dstack([img_np, np.full((h, w), 255, dtype=np.uint8)])
-    pixels_b64 = base64.b64encode(rgba.tobytes()).decode()
+    seg_map = run_segmentation(img, img_np, task_id=task_id)
+    pixels_b64 = base64.b64encode(contents).decode()
 
     primary_floor_ids = class_ids(PRIMARY_FLOOR_CLASSES)
     floor_class_ids = class_ids(FLOOR_SURFACE_CLASSES)
@@ -728,16 +789,16 @@ def build_segmentation_bundle(contents: bytes):
         floor_mask = wall_subtract(floor_mask, seg_map, dilation=1)
         floor_mask = clean_floor_mask(floor_mask)
 
-    depth = estimate_depth(img, w, h)
-    homography, plane = estimate_floor_plane(floor_mask, img_np)
+    depth = estimate_depth(img, w, h, task_id=task_id)
+    homography, plane = estimate_floor_plane(floor_mask, img_np, task_id=task_id)
     quad = np.asarray(plane["quad"], dtype=np.float32).reshape(4, 2) if plane and plane.get("quad") else None
-    surface_mask = build_floor_surface_mask(floor_mask, seg_map, quad, depth)
+    surface_mask = build_floor_surface_mask(floor_mask, seg_map, quad, depth, task_id=task_id)
     surface_indices = np.flatnonzero(surface_mask.ravel()).astype(np.uint32)
     shade_map = None
     if ENABLE_INTRINSIC_SHADING and intrinsic_models is not None:
-        shade_map = build_intrinsic_shade_map(img_np, surface_mask)
+        shade_map = build_intrinsic_shade_map(img_np, surface_mask, task_id=task_id)
     if shade_map is None:
-        shade_map = build_shade_map(img_np, surface_mask)
+        shade_map = build_shade_map(img_np, surface_mask, task_id=task_id)
     segments = []
 
     if len(surface_indices) >= min_floor_area:
@@ -777,6 +838,8 @@ def build_segmentation_bundle(contents: bytes):
                 },
             })
 
+    duration = time.perf_counter() - start_total
+    print(f"[{task_id}] Bundle build completed in {duration:.4f}s", flush=True)
     return {"width": w, "height": h, "pixels": pixels_b64, "segments": segments}
 
 
@@ -796,15 +859,20 @@ def write_job(job: dict):
 
 
 def run_conversion_task(job_id: str, upload_path: Path):
+    print(f"[{job_id}] Starting background conversion task...", flush=True)
+    start_time = time.perf_counter()
     try:
         image_bytes = upload_path.read_bytes()
-        bundle = build_segmentation_bundle(image_bytes)
+        bundle = build_segmentation_bundle(image_bytes, task_id=job_id)
         (JOB_DIR / f"{job_id}.bundle.json").write_text(json.dumps(bundle))
         job = read_job(job_id)
         job["status"] = "COMPLETED"
         write_job(job)
+        duration = time.perf_counter() - start_time
+        print(f"[{job_id}] Background conversion task completed in {duration:.4f}s", flush=True)
     except Exception as exc:
-        print(f"Background conversion failed: {exc}", flush=True)
+        duration = time.perf_counter() - start_time
+        print(f"[{job_id}] Background conversion failed after {duration:.4f}s: {exc}", flush=True)
         try:
             job = read_job(job_id)
             job["status"] = "FAILED"
@@ -866,7 +934,8 @@ async def viz2d_job_file(job_id: str):
 @app.post("/segment")
 async def segment(file: UploadFile = File(...)):
     contents = await file.read()
-    return build_segmentation_bundle(contents)
+    task_id = f"segment_{uuid.uuid4().hex[:8]}"
+    return build_segmentation_bundle(contents, task_id=task_id)
 
 
 if __name__ == "__main__":