crackmnist_eccv.run_experiment ¶

def evaluate(
    model: torch.nn.Module,
    loader: data.DataLoader,
    pollution_ratio: float,
    pollution_type: str,
    pollution_kwargs: dict,
) -> float:
    """Run model in eval mode and return average DiceLoss over the loader.

    Args:
        model: Trained PyTorch model.
        loader: DataLoader to evaluate on.
        pollution_ratio: Fraction of images to pollute.
        pollution_type: Type of pollution to apply.
        pollution_kwargs: Additional keyword arguments for the pollution
            function.

    Returns:
        Average DiceLoss as a float.
    """
    model.eval()
    criterion = DiceLoss()
    total_loss, total_count = 0.0, 0

    with torch.no_grad():
        for inputs, targets in loader:
            inputs = apply_per_image_pollution(
                inputs,
                pollution_ratio=pollution_ratio,
                pollution_type=pollution_type,
                **pollution_kwargs,
            )
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, targets.unsqueeze(1).float() / 255.0)
            batch_size = inputs.size(0)
            total_loss += loss.item() * batch_size
            total_count += batch_size

    return total_loss / total_count if total_count else 0.0

def fix_random_seeds(seed: int) -> None:
    """Fix all relevant random seeds for reproducibility.

    Args:
        seed: Integer seed value.
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.use_deterministic_algorithms(True)

def get_pollution_config(
    pollution_type: str, intensity: float | tuple[float, float]
) -> dict:
    """Return keyword arguments for apply_per_image_pollution.

    Args:
        pollution_type: One of 'blob', 'gaussian', 'salt_pepper',
            'gaussian_blur'.
        intensity: Float or (min, max) tuple depending on pollution type.

    Returns:
        Dict with 'pollution_type' and 'pollution_kwargs' keys.

    Raises:
        ValueError: If the pollution type is unsupported or intensity format
            is wrong.
    """
    if pollution_type == "blob":
        if not isinstance(intensity, tuple) or len(intensity) != 2:
            raise ValueError("Blob pollution requires intensity as a (min, max) tuple")
        return {
            "pollution_type": pollution_type,
            "pollution_kwargs": {
                "radius": RADIUS,
                "target_min_oar": intensity[0],
                "target_max_oar": intensity[1],
            },
        }
    elif pollution_type in ["gaussian", "salt_pepper", "gaussian_blur"]:
        if not isinstance(intensity, (float, int)):
            raise ValueError(f"{pollution_type} requires intensity as a float or int")
        return {
            "pollution_type": pollution_type,
            "pollution_kwargs": {"intensity": intensity},
        }
    else:
        raise ValueError(f"Unsupported pollution type: {pollution_type}")

def parse_python_literal(text: str, expected_type: type) -> list:
    """Convert a string like '[0.1,0.2]' or '[(0.1,0.2)]' to a Python object.

    Args:
        text: String representation of a Python literal.
        expected_type: The expected Python type (e.g. list).

    Returns:
        Parsed Python object of the expected type.

    Raises:
        ArgumentTypeError: If parsing fails or type does not match.
        TypeError: If the parsed value is not an instance of expected_type.
    """
    try:
        value = ast.literal_eval(text)
        if not isinstance(value, expected_type):
            raise TypeError
        return value
    except Exception as err:
        raise argparse.ArgumentTypeError(
            f"Value '{text}' is not a valid {expected_type.__name__}"
        ) from err

def train_and_evaluate(
    train_perc: float,
    test_perc: float,
    intensity: float | tuple[float, float],
    seed: int,
) -> tuple[float, float, float]:
    """Train a UNet and evaluate on val and test sets.

    Args:
        train_perc: Fraction of training images to pollute.
        test_perc: Fraction of test images to pollute during evaluation.
        intensity: Pollution intensity (float or tuple for blob).
        seed: Random seed for reproducibility.

    Returns:
        Tuple of (final_val_loss, final_test_loss, final_train_loss).
    """
    fix_random_seeds(seed)

    model = UNet(in_ch=2, out_ch=1, init_features=128, dropout_prob=0.2).to(device)
    criterion = DiceLoss()
    optimizer = optim.Adam(model.parameters(), lr=LR)

    config = get_pollution_config(SELECTED_POLLUTION_TYPE, intensity)
    pollution_type = config["pollution_type"]
    pollution_kwargs = config["pollution_kwargs"]

    model.train()
    for _epoch in range(NUM_EPOCHS):
        for inputs, targets in train_loader_full:
            inputs = apply_per_image_pollution(
                inputs,
                pollution_ratio=train_perc,
                pollution_type=pollution_type,
                **pollution_kwargs,
            )
            inputs, targets = inputs.to(device), targets.to(device)

            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, targets.unsqueeze(1).float() / 255.0)
            loss.backward()
            optimizer.step()

    final_val_loss = evaluate(
        model,
        val_loader_full,
        pollution_ratio=train_perc,
        pollution_type=pollution_type,
        pollution_kwargs=pollution_kwargs,
    )
    final_test_loss = evaluate(
        model,
        test_loader_full,
        pollution_ratio=test_perc,
        pollution_type=pollution_type,
        pollution_kwargs=pollution_kwargs,
    )
    final_train_loss = evaluate(
        model,
        train_loader_full,
        pollution_ratio=train_perc,
        pollution_type=pollution_type,
        pollution_kwargs=pollution_kwargs,
    )

    return final_val_loss, final_test_loss, final_train_loss