CYWS-3D Model: Forward Pass, Predict, and Loss API

The Model class is the central PyTorch module of CYWS-3D. It accepts a prepared batch dict, computes paired CenterNet outputs for both images in a scene pair, and decodes them into bounding boxes that mark changed regions. The model is configured entirely through an EasyDict loaded from config.yml and optionally initialised from a .ckpt checkpoint file. All public methods follow standard nn.Module conventions.

Constructor

from modules.model import Model
from utils import get_easy_dict_from_yaml_file

args = get_easy_dict_from_yaml_file("config.yml")
model = Model(args, load_weights_from="checkpoints/cyws3d.ckpt")

args

EasyDict

required

Configuration object produced by get_easy_dict_from_yaml_file("config.yml"). Controls model architecture choices such as backbone, CenterNet head dimensions, and registration module settings. The structure mirrors the top-level keys in config.yml.

load_weights_from

str

default:"None"

Path to a .ckpt checkpoint file. When supplied, the constructor calls safely_load_state_dict to restore weights. When None, the model starts with random initialisation — only useful for training from scratch.

Model.predict(batch)

batch_image1_bboxes, batch_image2_bboxes = model.predict(batch)

Runs the full forward pass under torch.no_grad() and decodes the raw CenterNet logits into bounding boxes via get_bboxes_from_logits. This is the recommended entry point for inference. Returns a tuple of two lists, each of length batch_size:

Return value	Type	Description
`batch_image1_predicted_bboxes`	`List[Tuple[Tensor, Tensor]]`	Per-image-1 detections for every batch item
`batch_image2_predicted_bboxes`	`List[Tuple[Tensor, Tensor]]`	Per-image-2 detections for every batch item

Each element of those lists is a (bboxes_tensor, labels_tensor) tuple where:

bboxes_tensor — shape [N, 5], dtype float. The last dimension is [x1, y1, x2, y2, score] in 224×224 pixel coordinates.
labels_tensor — shape [N], dtype long. Class labels for each detection (single-class models always emit 0).

predict is decorated with @torch.no_grad() — you do not need to wrap the call in a with torch.no_grad() block yourself.

Post-process the raw tensors with the helpers from utils.py before visualisation:

import numpy as np
from utils import (
    remove_bboxes_with_area_less_than,
    suppress_overlapping_bboxes,
    keep_matching_bboxes,
)

bboxes1, bboxes2 = batch_image1_bboxes[0][0].cpu().numpy(), batch_image2_bboxes[0][0].cpu().numpy()
bboxes1 = remove_bboxes_with_area_less_than(bboxes1, area_threshold=400)
bboxes2 = remove_bboxes_with_area_less_than(bboxes2, area_threshold=400)
bboxes1, scores1 = suppress_overlapping_bboxes(bboxes1[:, :4], bboxes1[:, 4])
bboxes2, scores2 = suppress_overlapping_bboxes(bboxes2[:, :4], bboxes2[:, 4])

Model.forward(batch)

image1_outputs, image2_outputs = model(batch)

Runs a single forward pass and returns the raw CenterNet head outputs for both images. Use this method when you need access to intermediate logits — for example, during training or custom evaluation. Returns a tuple of two tuples:

Return value	Type	Description
`image1_centernet_outputs`	`Tuple[Tensor, Tensor, Tensor]`	`(heatmap, wh, offset)` for image 1
`image2_centernet_outputs`	`Tuple[Tensor, Tensor, Tensor]`	`(heatmap, wh, offset)` for image 2

Each output triple follows the standard CenterNet convention:

heatmap — class probability map, shape [B, C, H', W']
wh — bounding box width/height predictions, shape [B, 2, H', W']
offset — sub-pixel offset correction, shape [B, 2, H', W']

Model.compute_loss(batch, image1_outputs, image2_outputs)

loss = model.compute_loss(batch, image1_outputs, image2_outputs)
loss.backward()

Computes the combined CenterNet detection loss for both images and returns a scalar Tensor. The loss terms for image 1 and image 2 are summed. The batch dict must contain the following training-only keys in addition to the standard inference keys:

Key	Type	Description
`target_bbox_1`	`Tensor`	Ground-truth bounding boxes for image 1
`target_bbox_labels1`	`Tensor`	Class labels for each ground-truth box in image 1
`target_bbox_2`	`Tensor`	Ground-truth bounding boxes for image 2
`target_bbox_labels2`	`Tensor`	Class labels for each ground-truth box in image 2
`query_metadata`	`dict`	CenterNet head metadata (stride, output size); added by `prepare_batch_for_model`

compute_loss does not call forward internally. You must pass the outputs from a preceding model(batch) call. This design allows you to reuse the same forward pass for both loss computation and logging intermediate activations.

Model.safely_load_state_dict(checkpoint_state_dict)

model.safely_load_state_dict(torch.load("checkpoints/cyws3d.ckpt")["state_dict"])

A fault-tolerant wrapper around the standard load_state_dict. Instead of raising an error on shape mismatches, it:

Skips any parameter whose checkpoint shape differs from the current model shape, keeping the randomly initialised value.
Drops checkpoint keys that do not exist in the current model (e.g. keys left over from a previous architecture version).

This makes it safe to fine-tune a checkpoint on a model whose head dimensions have been modified, or to load weights from a checkpoint that was saved with a slightly different configuration.

The constructor calls safely_load_state_dict automatically when load_weights_from is provided. Call it directly only when you need to swap weights after construction.

Batch Format

Model.forward and Model.predict expect a batch dictionary with the following keys. The batch is assembled by create_batch_from_metadata and finalised by prepare_batch_for_model.

Key	Shape / Type	Description
`image1`	`Tensor [B, C, 224, 224]`	ImageNet-normalised RGB image 1, stacked across the batch
`image2`	`Tensor [B, C, 224, 224]`	ImageNet-normalised RGB image 2, stacked across the batch
`depth1`	per-item `Tensor`	Metric or relative depth map for image 1; predicted by ZoeDepth if absent in metadata
`depth2`	per-item `Tensor`	Metric or relative depth map for image 2
`intrinsics1`	per-item `ndarray (3,3)`	Camera intrinsics matrix for image 1; adjusted for 224×224 resize
`intrinsics2`	per-item `ndarray (3,3)`	Camera intrinsics matrix for image 2
`position1`	per-item `ndarray (3,)`	World-space translation of camera 1
`position2`	per-item `ndarray (3,)`	World-space translation of camera 2
`rotation1`	per-item `ndarray (3,3)`	Rotation matrix for camera 1
`rotation2`	per-item `ndarray (3,3)`	Rotation matrix for camera 2
`transfm2d_1_to_2`	per-item `ndarray (3,3)`	2-D homography from image 1 to image 2
`transfm2d_2_to_1`	per-item `ndarray (3,3)`	2-D homography from image 2 to image 1
`registration_strategy`	`List[str]`	One entry per batch item: `"3d"`, `"2d"`, `"2d_from_corr"`, or `"identity"`
`query_metadata`	`dict`	CenterNet stride and output resolution metadata; added by `prepare_batch_for_model`
`points1`	per-item `Tensor`	Normalised `[0,1]` keypoint coordinates in image 1; added by `CorrespondenceExtractor`
`points2`	per-item `Tensor`	Normalised `[0,1]` keypoint coordinates in image 2; added by `CorrespondenceExtractor`

Keys related to camera geometry (intrinsics, position, rotation, transfm2d_*) are only required when the corresponding registration_strategy needs them. The model reads registration_strategy per item and routes accordingly.

Full Inference Example

import torch
from modules.model import Model
from modules.correspondence_extractor import CorrespondenceExtractor
from utils import (
    get_easy_dict_from_yaml_file,
    create_batch_from_metadata,
    fill_in_the_missing_information,
    prepare_batch_for_model,
)

# 1. Load config and build model
configs = get_easy_dict_from_yaml_file("config.yml")
model = Model(configs, load_weights_from="checkpoints/cyws3d.ckpt")
model.eval()

# 2. Build auxiliary components
correspondence_extractor = CorrespondenceExtractor()
depth_predictor = torch.hub.load(
    "isl-org/ZoeDepth", "ZoeD_NK", pretrained=True
).eval()

# 3. Assemble batch
batch_metadata = get_easy_dict_from_yaml_file("demo_data/input_metadata.yml")
batch = create_batch_from_metadata(batch_metadata)
batch = fill_in_the_missing_information(batch, depth_predictor, correspondence_extractor)
batch = prepare_batch_for_model(batch)

# 4. Run prediction
batch_bboxes1, batch_bboxes2 = model.predict(batch)

# 5. Inspect first item
bboxes1, labels1 = batch_bboxes1[0]
print(f"Detected {bboxes1.shape[0]} changes in image 1")
print(f"Top box (x1,y1,x2,y2,score): {bboxes1[0].cpu().numpy()}")

Get Started

Concepts

Datasets

Reference

CYWS-3D Model: Forward Pass, Predict, and Loss API

Constructor

Model.predict(batch)

Model.forward(batch)

Model.compute_loss(batch, image1_outputs, image2_outputs)

Model.safely_load_state_dict(checkpoint_state_dict)

Batch Format

Full Inference Example

Build docs developers (and LLMs) love

Get Started

Concepts

Datasets

Reference

Documentation Index

​Constructor

​Model.predict(batch)

​Model.forward(batch)

​Model.compute_loss(batch, image1_outputs, image2_outputs)

​Model.safely_load_state_dict(checkpoint_state_dict)

​Batch Format

​Full Inference Example

Build docs developers (and LLMs) love

Constructor

Model.predict(batch)

Model.forward(batch)

Model.compute_loss(batch, image1_outputs, image2_outputs)

Model.safely_load_state_dict(checkpoint_state_dict)

Batch Format

Full Inference Example