zamba.images.manager

get_default_transforms(model_family, image_size=None)

Build the (top, bottom) eval transform lists for a preprocessing family.

The preprocessing is fully determined by model_family (and the resolved image_size), NOT by a model_name string on the config. This lets prediction derive transforms directly from a loaded checkpoint. Returns the transform lists plus the resolved integer image size.

Source code in zamba/images/manager.py

def get_default_transforms(model_family: str, image_size=None):
    """Build the (top, bottom) eval transform lists for a preprocessing family.

    The preprocessing is fully determined by ``model_family`` (and the resolved
    ``image_size``), NOT by a model_name string on the config. This lets prediction
    derive transforms directly from a loaded checkpoint. Returns the transform lists
    plus the resolved integer image size.
    """
    # checkpoints may store image_size as a tuple (e.g. (480, 480)); normalize to int
    if isinstance(image_size, (tuple, list)):
        image_size = image_size[0]

    logger.info(f"Using default transforms for '{model_family}' model family")
    if model_family == ImageModelEnum.SPECIESNET.value:
        # speciesnet is trained on 480x480 images by default
        if image_size is None:
            logger.info("Image size not specified, using value from model checkpoint: 480")
            image_size = 480

        top_transforms = [
            transforms.Resize(
                (image_size, image_size),
                interpolation=transforms.InterpolationMode.BICUBIC,
            ),
        ]
        bottom_transforms = [
            transforms.ToTensor(),
        ]
    else:
        # lila.science and generic models: pad to square + ImageNet-ish normalization
        if image_size is None:
            logger.info("Image size not specified, using default value: 224")
            image_size = 224

        top_transforms = [
            transforms.Lambda(partial(resize_and_pad, desired_size=(image_size, image_size))),
        ]
        bottom_transforms = [
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.45, 0.45, 0.45], std=[0.225, 0.225, 0.225]),
        ]

    return top_transforms, bottom_transforms, image_size

resolve_inference_family(model_name, checkpoint)

Determine the preprocessing family without trusting a (possibly stale/mangled) model_name default.

When a checkpoint is provided it is authoritative: the family is read from the checkpoint's persisted model_family / legacy zamba_model hparam, falling back to inference from the stored architecture name. Only when there is no checkpoint do we fall back to the configured model_name.

Source code in zamba/images/manager.py

def resolve_inference_family(model_name, checkpoint) -> str:
    """Determine the preprocessing family without trusting a (possibly stale/mangled)
    ``model_name`` default.

    When a checkpoint is provided it is authoritative: the family is read from the
    checkpoint's persisted ``model_family`` / legacy ``zamba_model`` hparam, falling
    back to inference from the stored architecture name. Only when there is no
    checkpoint do we fall back to the configured ``model_name``.
    """
    if checkpoint is not None:
        try:
            hp = get_checkpoint_hparams(checkpoint)
            family = hp.get("model_family") or hp.get("zamba_model")
            if family:
                return family
            return infer_model_family(hp.get("model_name"))
        except Exception as exc:  # noqa: BLE001 -- fall back to model_name on any read error
            logger.warning(f"Could not read family from checkpoint ({exc}); using model_name.")
    return infer_model_family(model_name)

resolve_training_image_size(config)

Resolve the image size to train at.

An explicitly configured image_size always wins. Otherwise, when finetuning or resuming from a checkpoint, the checkpoint's own image_size takes precedence over the preprocessing-family default, since a finetuned model may have been trained at a non-default size. Returns None (deferring to the family default) only when there is no explicit size and no usable size on the checkpoint. The returned value may be a scalar or a tuple; get_default_transforms normalizes it to an int.

Source code in zamba/images/manager.py

def resolve_training_image_size(config: ImageClassificationTrainingConfig):
    """Resolve the image size to train at.

    An explicitly configured ``image_size`` always wins. Otherwise, when finetuning or
    resuming from a checkpoint, the checkpoint's own ``image_size`` takes precedence over
    the preprocessing-family default, since a finetuned model may have been trained at a
    non-default size. Returns ``None`` (deferring to the family default) only when there is
    no explicit size and no usable size on the checkpoint. The returned value may be a
    scalar or a tuple; ``get_default_transforms`` normalizes it to an int.
    """
    if config.image_size is not None:
        return config.image_size

    if config.checkpoint is not None and not config.from_scratch:
        try:
            return get_checkpoint_hparams(config.checkpoint).get("image_size")
        except Exception as exc:  # noqa: BLE001 -- fall back to the family default
            logger.warning(f"Could not read image_size from checkpoint ({exc}); using default.")

    return None