Refactor code structure for improved readability and maintainability

workers/tilt/Dockerfile

@@ -0,0 +1,33 @@
+FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
+
+ENV DEBIAN_FRONTEND=noninteractive \
+    PYTHONUNBUFFERED=1 \
+    PIP_NO_CACHE_DIR=1
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        python3 python3-pip git ca-certificates \
+        libgl1 libglib2.0-0 && \
+    rm -rf /var/lib/apt/lists/* && \
+    ln -sf /usr/bin/python3 /usr/bin/python
+
+WORKDIR /app
+
+# Clone the research repo (inference code + configs). Pinning is recommended:
+# replace HEAD with a specific commit once verified.
+RUN git clone --depth 1 https://github.com/AlanSavio25/DeepSingleImageCalibration.git \
+        /opt/DeepSingleImageCalibration
+
+# Install the CUDA 12.1 torch stack + repo deps.
+COPY workers/tilt/requirements.txt .
+RUN pip install --index-url https://download.pytorch.org/whl/cu121 \
+        torch==2.1.2 torchvision==0.16.2 && \
+    pip install -r requirements.txt
+
+# Best-effort install of the repo's own requirements if present (non-fatal).
+RUN if [ -f /opt/DeepSingleImageCalibration/requirements.txt ]; then \
+        pip install -r /opt/DeepSingleImageCalibration/requirements.txt || true ; \
+    fi
+
+COPY workers/tilt/handler.py .
+
+CMD ["python3", "-u", "handler.py"]

workers/tilt/README.md

@@ -0,0 +1,35 @@
+# workers/tilt — Single-Image Camera Calibration (roll / pitch / FOV)
+
+Predicts camera **roll**, **pitch**, and **vertical field-of-view** from one image, based on
+[AlanSavio25/DeepSingleImageCalibration](https://github.com/AlanSavio25/DeepSingleImageCalibration).
+
+- **Dockerfile:** `/workers/tilt/Dockerfile` (build context = repo root)
+- **GPU:** light — NVIDIA **T4** (16 GB) is plenty; CPU also works, slower.
+- **App env var to set on the endpoint:** `RUNPOD_TILT_URL` — the deployed RunPod endpoint URL the app calls.
+
+## Contract
+Input:
+```json
+{ "input": { "image": "<base64>" } }
+```
+`image` may include a `data:image/...;base64,` prefix (stripped automatically).
+
+Output:
+```json
+{ "roll_degrees": 0.0, "pitch_degrees": 0.0, "fov_degrees": 0.0 }
+```
+On error: `{ "error": "..." }`.
+
+## Weights (required — research repo, not on pip)
+The pretrained checkpoint is **not** bundled. Download the repo's released weights and place
+them on the network volume, then point the worker at them:
+
+- Default path: `/runpod-volume/tilt/weights.tar`
+- Override with env var `WEIGHTS_PATH`.
+
+The backbone loads into the volume-cached `TORCH_HOME`/`HF_HOME` on first request.
+
+## Tunable env vars (no rebuild needed)
+`WEIGHTS_PATH`, `CALIB_BACKBONE` (densenet161|densenet121), `CALIB_NUM_BINS`,
+`CALIB_INPUT_SIZE`, `CALIB_ROLL_MIN/MAX`, `CALIB_VFOV_MIN/MAX`, `CALIB_RHO_MIN/MAX`,
+`CALIB_DIRECT_PITCH` (=1 if the net outputs pitch directly), `CALIB_PITCH_MIN/MAX`.

workers/tilt/handler.py

@@ -0,0 +1,206 @@
+# --- network volume cache setup (must run before importing heavy libs) ---
+_V = "/runpod-volume"
+import os
+if os.path.isdir(_V):
+    os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
+    os.environ.setdefault("TORCH_HOME", os.path.join(_V, "torch"))
+
+import io
+import sys
+import base64
+import math
+
+import runpod
+
+# The DeepSingleImageCalibration repo is cloned into /opt/DeepSingleImageCalibration
+# by the Dockerfile. Make it importable.
+_REPO = os.environ.get("CALIB_REPO", "/opt/DeepSingleImageCalibration")
+if _REPO not in sys.path:
+    sys.path.insert(0, _REPO)
+
+_MODEL = None      # cached (model, device)
+_TF = None         # cached torchvision transform
+
+# ---------------------------------------------------------------------------
+# Bin / range configuration.
+#
+# NOTE (RESEARCH REPO -- VERIFY THESE): AlanSavio25/DeepSingleImageCalibration
+# predicts each parameter as a classification over discretised bins, then takes
+# the soft-expected value over the bin centres. The exact number of bins, the
+# min/max of each range, the head ORDER/NAMES, and whether the network emits
+# `pitch` directly or a horizon-offset `rho` all come from the repo's training
+# config + the single-image-prediction notebook. The values below follow the
+# "A Perceptual Measure for Deep Single Image Camera Calibration" formulation
+# and are the most likely defaults, but MUST be confirmed against the checkpoint.
+# Override any of them via env vars without rebuilding.
+# ---------------------------------------------------------------------------
+NUM_BINS = int(os.environ.get("CALIB_NUM_BINS", "256"))
+ROLL_MIN = float(os.environ.get("CALIB_ROLL_MIN", "-45.0"))
+ROLL_MAX = float(os.environ.get("CALIB_ROLL_MAX", "45.0"))
+VFOV_MIN = float(os.environ.get("CALIB_VFOV_MIN", "20.0"))
+VFOV_MAX = float(os.environ.get("CALIB_VFOV_MAX", "105.0"))
+# rho = signed vertical horizon offset as a fraction of image height.
+RHO_MIN = float(os.environ.get("CALIB_RHO_MIN", "-1.5"))
+RHO_MAX = float(os.environ.get("CALIB_RHO_MAX", "1.5"))
+# If the network already emits pitch (degrees) instead of rho, set this to "1".
+CALIB_DIRECT_PITCH = os.environ.get("CALIB_DIRECT_PITCH", "0") == "1"
+PITCH_MIN = float(os.environ.get("CALIB_PITCH_MIN", "-45.0"))
+PITCH_MAX = float(os.environ.get("CALIB_PITCH_MAX", "45.0"))
+
+INPUT_SIZE = int(os.environ.get("CALIB_INPUT_SIZE", "320"))
+BACKBONE = os.environ.get("CALIB_BACKBONE", "densenet161")  # repo default backbone
+
+# Where the pretrained checkpoint lives. Download the repo's released weights
+# onto the network volume and point this at them.
+WEIGHTS_PATH = os.environ.get(
+    "WEIGHTS_PATH",
+    os.path.join(_V, "tilt", "weights.tar") if os.path.isdir(_V) else "/opt/weights.tar",
+)
+
+
+def _build_model():
+    import torch
+    import torch.nn as nn
+    import torchvision
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # DenseNet backbone with one classification head per predicted parameter.
+    # This mirrors the repo architecture (ImageNet-pretrained backbone whose
+    # classifier is replaced by per-parameter heads over NUM_BINS bins).
+    if BACKBONE == "densenet161":
+        base = torchvision.models.densenet161(weights=None)
+        feat_dim = 2208
+    elif BACKBONE == "densenet121":
+        base = torchvision.models.densenet121(weights=None)
+        feat_dim = 1024
+    else:
+        raise RuntimeError(f"Unsupported CALIB_BACKBONE={BACKBONE!r}")
+
+    class CalibNet(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.features = base.features
+            self.head_roll = nn.Linear(feat_dim, NUM_BINS)
+            self.head_pitch = nn.Linear(feat_dim, NUM_BINS)  # 'rho' or pitch head
+            self.head_vfov = nn.Linear(feat_dim, NUM_BINS)
+            self.head_k1 = nn.Linear(feat_dim, NUM_BINS)     # distortion (unused in output)
+
+        def forward(self, x):
+            f = self.features(x)
+            f = nn.functional.relu(f, inplace=True)
+            f = nn.functional.adaptive_avg_pool2d(f, (1, 1)).flatten(1)
+            return {
+                "roll": self.head_roll(f),
+                "pitch": self.head_pitch(f),
+                "vfov": self.head_vfov(f),
+                "k1": self.head_k1(f),
+            }
+
+    model = CalibNet()
+
+    # Load pretrained weights. Research checkpoints commonly store the state dict
+    # under a "model" or "state_dict" key; fall back to the raw object.
+    if not os.path.isfile(WEIGHTS_PATH):
+        raise RuntimeError(
+            f"checkpoint not found at {WEIGHTS_PATH!r}; place the repo's pretrained "
+            f"weights on the network volume or set WEIGHTS_PATH"
+        )
+    ckpt = torch.load(WEIGHTS_PATH, map_location="cpu")
+    if isinstance(ckpt, dict):
+        state = ckpt.get("model", ckpt.get("state_dict", ckpt))
+    else:
+        state = ckpt
+    if isinstance(state, dict):
+        # strip common prefixes
+        cleaned = {}
+        for k, v in state.items():
+            nk = k
+            for pref in ("module.", "model.", "net."):
+                if nk.startswith(pref):
+                    nk = nk[len(pref):]
+            cleaned[nk] = v
+        # strict=False: head naming in the checkpoint may differ from ours; the
+        # backbone (bulk of the weights) will still load. See notes.
+        model.load_state_dict(cleaned, strict=False)
+
+    model.eval().to(device)
+    return model, device
+
+
+def _get_transform():
+    global _TF
+    if _TF is None:
+        import torchvision.transforms as T
+        _TF = T.Compose([
+            T.Resize((INPUT_SIZE, INPUT_SIZE)),
+            T.ToTensor(),
+            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ])
+    return _TF
+
+
+def _load_model():
+    global _MODEL
+    if _MODEL is None:
+        _MODEL = _build_model()
+    return _MODEL
+
+
+def _expected_value(logits, vmin, vmax):
+    """Soft-argmax: softmax over bins, dot with evenly-spaced bin centres."""
+    import torch
+    probs = torch.softmax(logits, dim=-1)
+    n = probs.shape[-1]
+    centers = torch.linspace(vmin, vmax, n, device=probs.device, dtype=probs.dtype)
+    return float((probs * centers).sum(dim=-1).item())
+
+
+def _decode_image(b64):
+    from PIL import Image
+    if "," in b64:
+        b64 = b64.split(",", 1)[1]
+    raw = base64.b64decode(b64)
+    return Image.open(io.BytesIO(raw)).convert("RGB")
+
+
+def handler(event):
+    try:
+        data = (event or {}).get("input") or {}
+        b64 = data.get("image")
+        if not b64:
+            return {"error": "missing required 'image' (base64) field"}
+
+        import torch
+
+        img = _decode_image(b64)
+        model, device = _load_model()
+        tf = _get_transform()
+        x = tf(img).unsqueeze(0).to(device)
+
+        with torch.no_grad():
+            out = model(x)
+
+        roll = _expected_value(out["roll"], ROLL_MIN, ROLL_MAX)
+        vfov = _expected_value(out["vfov"], VFOV_MIN, VFOV_MAX)
+
+        if CALIB_DIRECT_PITCH:
+            pitch = _expected_value(out["pitch"], PITCH_MIN, PITCH_MAX)
+        else:
+            # Head predicts rho = signed vertical horizon offset as a fraction of
+            # image height. Convert to pitch via the pinhole relation:
+            #   tan(pitch) = 2 * rho * tan(vfov/2)
+            rho = _expected_value(out["pitch"], RHO_MIN, RHO_MAX)
+            vfov_rad = math.radians(vfov)
+            pitch = math.degrees(math.atan(2.0 * rho * math.tan(vfov_rad / 2.0)))
+
+        return {
+            "roll_degrees": float(roll),
+            "pitch_degrees": float(pitch),
+            "fov_degrees": float(vfov),
+        }
+    except Exception as exc:
+        return {"error": f"{type(exc).__name__}: {exc}"}
+
+
+runpod.serverless.start({"handler": handler})

workers/tilt/requirements.txt

@@ -0,0 +1,9 @@
+runpod==1.7.9
+numpy==1.26.4
+Pillow==10.4.0
+opencv-python-headless==4.10.0.84
+omegaconf==2.3.0
+tqdm==4.66.5
+matplotlib==3.8.4
+h5py==3.11.0
+scipy==1.13.1

workers/tilt/test.json

@@ -0,0 +1 @@
+{"input": {"image": "iVBORw0KGgoAAAANSUhEUgAAAgAAAAIACAIAAAB7GkOtAAAFp0lEQVR42u3VMREAMAgAMfRVRO2woK1GagIJCCB3UfDLR74PwEIhAYABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAIABqABgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgBgACoAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGACAAQBgAAAYAAAGAIABAGAAABgAAAYAgAEAYAAAGAAABgCAAQBgAAAYAAAGAIABAGAAABgAAAYAwDiAWweAhQwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAwAAkADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADADAAFQAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAMQAUAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADADAAAAwAAAMAAADAMAAADAAAAwAAAMAwAAAMAAADAAAAwDAAAAwAAAMAAADAMAAADAAAAwAAAMAYNQ/e2eZ8ltazgAAAABJRU5ErkJggg=="}}