207 lines
7.5 KiB
Python
207 lines
7.5 KiB
Python
# --- 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})
|