visionkit-pro 0.1.2__py3-none-any.whl

visionkit/__init__.py

@@ -0,0 +1,23 @@
+"""Reusable image-classification training, evaluation, and Grad-CAM tools."""
+
+from visionkit.config import ExperimentConfig
+from visionkit.data import CSVImageDataset, build_dataloaders, build_transforms
+from visionkit.gradcam import GradCAM, generate_gradcam_report
+from visionkit.models import build_model, load_checkpoint
+from visionkit.pipeline import run_evaluation, run_gradcam, run_training
+from visionkit.train import train_model
+
+__all__ = [
+    "CSVImageDataset",
+    "ExperimentConfig",
+    "GradCAM",
+    "build_dataloaders",
+    "build_model",
+    "build_transforms",
+    "generate_gradcam_report",
+    "load_checkpoint",
+    "run_evaluation",
+    "run_gradcam",
+    "run_training",
+    "train_model",
+]

visionkit/cli.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from typing import Optional
+
+from visionkit.config import ExperimentConfig
+from visionkit.pipeline import run_evaluation, run_gradcam, run_training
+
+
+def _image_size(value: str) -> tuple[int, int]:
+    parts = value.lower().replace("x", ",").split(",")
+    if len(parts) != 2:
+        raise argparse.ArgumentTypeError("image size must look like 300x300")
+    return int(parts[0]), int(parts[1])
+
+
+def _base_parser(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument("--csv", dest="csv_path", type=Path, required=True)
+    parser.add_argument("--image-dir", type=Path, required=True)
+    parser.add_argument("--output-dir", type=Path, required=True)
+    parser.add_argument("--filename-column", default="filename")
+    parser.add_argument("--label-column", default="label")
+    parser.add_argument("--split-column", default="split")
+    parser.add_argument("--classes", nargs="*")
+    parser.add_argument("--architecture", default="efficientnet_b3")
+    parser.add_argument("--weights", default="DEFAULT")
+    parser.add_argument("--image-size", type=_image_size, default=(300, 300))
+    parser.add_argument("--batch-size", type=int, default=8)
+    parser.add_argument("--seed", type=int, default=123)
+    parser.add_argument("--positive-class")
+    parser.add_argument("--target-layer")
+    parser.add_argument("--device")
+
+
+def _config(args) -> ExperimentConfig:
+    return ExperimentConfig(
+        csv_path=args.csv_path,
+        image_dir=args.image_dir,
+        output_dir=args.output_dir,
+        filename_column=args.filename_column,
+        label_column=args.label_column,
+        split_column=args.split_column,
+        train_split=getattr(args, "train_split", "train"),
+        val_split=getattr(args, "val_split", "validation"),
+        class_names=args.classes,
+        architecture=args.architecture,
+        weights=args.weights,
+        image_size=args.image_size,
+        batch_size=args.batch_size,
+        epochs=getattr(args, "epochs", 30),
+        learning_rate=getattr(args, "learning_rate", 1e-3),
+        seed=args.seed,
+        positive_class=args.positive_class,
+        target_layer=args.target_layer,
+        device=args.device,
+    )
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog="visionkit")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    train = subparsers.add_parser("train")
+    _base_parser(train)
+    train.add_argument("--train-split", default="train")
+    train.add_argument("--val-split", default="validation")
+    train.add_argument("--epochs", type=int, default=30)
+    train.add_argument("--learning-rate", type=float, default=1e-3)
+
+    evaluate = subparsers.add_parser("evaluate")
+    _base_parser(evaluate)
+    evaluate.add_argument("--checkpoint", type=Path, required=True)
+    evaluate.add_argument("--split", default="validation")
+    evaluate.add_argument("--output-csv", type=Path)
+
+    gradcam = subparsers.add_parser("gradcam")
+    _base_parser(gradcam)
+    gradcam.add_argument("--checkpoint", type=Path, required=True)
+    gradcam.add_argument("--split", default="validation")
+    gradcam.add_argument("--gradcam-output-dir", type=Path)
+    gradcam.add_argument("--output-csv", type=Path)
+    gradcam.add_argument("--target-class")
+    return parser
+
+
+def main(argv: Optional[list[str]] = None) -> None:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    config = _config(args)
+
+    if args.command == "train":
+        _, history = run_training(config)
+        print(history.tail(1).to_string(index=False))
+    elif args.command == "evaluate":
+        _, metrics = run_evaluation(config, args.checkpoint, split_name=args.split, output_csv=args.output_csv)
+        print(metrics)
+    elif args.command == "gradcam":
+        _, metrics = run_gradcam(
+            config,
+            args.checkpoint,
+            split_name=args.split,
+            output_dir=args.gradcam_output_dir,
+            output_csv=args.output_csv,
+            target_class=args.target_class,
+        )
+        print(metrics)
+
+
+if __name__ == "__main__":
+    main()

visionkit/config.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Optional, Sequence
+
+
+@dataclass
+class ExperimentConfig:
+    csv_path: Path
+    image_dir: Path
+    output_dir: Path
+    filename_column: str = "filename"
+    label_column: str = "label"
+    split_column: Optional[str] = "split"
+    train_split: str = "train"
+    val_split: str = "validation"
+    class_names: Optional[Sequence[str]] = None
+    architecture: str = "efficientnet_b3"
+    weights: Optional[str] = "DEFAULT"
+    image_size: tuple[int, int] = (300, 300)
+    batch_size: int = 8
+    epochs: int = 30
+    learning_rate: float = 1e-3
+    num_workers: Optional[int] = None
+    seed: int = 123
+    positive_class: Optional[str] = None
+    target_layer: Optional[str] = None
+    device: Optional[str] = None
+    checkpoint_metric: str = "val_auc"
+    save_every_epoch: bool = True
+    normalize_mean: tuple[float, float, float] = (0.485, 0.456, 0.406)
+    normalize_std: tuple[float, float, float] = (0.229, 0.224, 0.225)
+    extra_metadata_columns: Sequence[str] = field(default_factory=tuple)
+
+    def model_dir(self) -> Path:
+        return self.output_dir / "models"
+
+    def log_dir(self) -> Path:
+        return self.output_dir / "logs"
+
+    def prediction_dir(self) -> Path:
+        return self.output_dir / "predictions"
+
+    def gradcam_dir(self) -> Path:
+        return self.output_dir / "gradcam"
+
+    def ensure_output_dirs(self) -> None:
+        for path in [
+            self.output_dir,
+            self.model_dir(),
+            self.log_dir(),
+            self.prediction_dir(),
+            self.gradcam_dir(),
+        ]:
+            path.mkdir(parents=True, exist_ok=True)

visionkit/data.py

@@ -0,0 +1,175 @@
+from __future__ import annotations
+
+from functools import partial
+from pathlib import Path
+from typing import Optional, Sequence
+
+import pandas as pd
+import torch
+from PIL import Image
+from torch.utils.data import DataLoader, Dataset
+from torchvision import transforms
+
+from visionkit.config import ExperimentConfig
+from visionkit.reproducibility import seed_worker
+
+
+class CSVImageDataset(Dataset):
+    def __init__(
+        self,
+        rows: pd.DataFrame,
+        image_dir: Path,
+        filename_column: str = "filename",
+        label_column: Optional[str] = "label",
+        class_to_idx: Optional[dict[str, int]] = None,
+        transform=None,
+        return_metadata: bool = False,
+    ) -> None:
+        self.rows = rows.reset_index(drop=True).copy()
+        self.image_dir = Path(image_dir)
+        self.filename_column = filename_column
+        self.label_column = label_column
+        self.class_to_idx = class_to_idx or {}
+        self.transform = transform
+        self.return_metadata = return_metadata
+        self.samples: list[tuple[Path, Optional[int], dict]] = []
+
+        missing_files = []
+        unknown_labels = []
+        for record in self.rows.to_dict(orient="records"):
+            filename = str(record[filename_column])
+            image_path = self.image_dir / filename
+            if not image_path.exists():
+                missing_files.append(filename)
+                continue
+
+            target = None
+            if label_column is not None and label_column in record:
+                label = str(record[label_column])
+                if label not in self.class_to_idx:
+                    unknown_labels.append(label)
+                    continue
+                target = self.class_to_idx[label]
+            self.samples.append((image_path, target, record))
+
+        if missing_files:
+            preview = ", ".join(missing_files[:5])
+            raise FileNotFoundError(f"Images listed in CSV were not found: {preview}")
+        if unknown_labels:
+            labels = ", ".join(sorted(set(unknown_labels)))
+            raise ValueError(f"Labels are missing from class mapping: {labels}")
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, index: int):
+        image_path, target, record = self.samples[index]
+        image = Image.open(image_path).convert("RGB")
+        if self.transform is not None:
+            image = self.transform(image)
+        if self.return_metadata:
+            return image, target, record
+        return image, target
+
+
+def read_metadata(csv_path: Path) -> pd.DataFrame:
+    return pd.read_csv(csv_path)
+
+
+def infer_class_names(df: pd.DataFrame, label_column: str, class_names: Optional[Sequence[str]]) -> list[str]:
+    if class_names:
+        return [str(name) for name in class_names]
+    return sorted(str(label) for label in df[label_column].dropna().unique())
+
+
+def class_mapping(class_names: Sequence[str]) -> tuple[dict[str, int], list[str]]:
+    names = [str(name) for name in class_names]
+    return {name: idx for idx, name in enumerate(names)}, names
+
+
+def filter_split(df: pd.DataFrame, split_column: Optional[str], split_name: Optional[str]) -> pd.DataFrame:
+    if not split_column or split_name is None:
+        return df.copy()
+    if split_column not in df.columns:
+        raise KeyError(f"Split column '{split_column}' was not found in CSV.")
+    return df[df[split_column].astype(str) == str(split_name)].copy()
+
+
+def build_transforms(
+    image_size: tuple[int, int] = (300, 300),
+    train: bool = False,
+    mean: tuple[float, float, float] = (0.485, 0.456, 0.406),
+    std: tuple[float, float, float] = (0.229, 0.224, 0.225),
+):
+    ops = [transforms.Resize(image_size)]
+    if train:
+        ops.extend(
+            [
+                transforms.RandomHorizontalFlip(),
+                transforms.RandomVerticalFlip(),
+                transforms.RandomRotation(10),
+            ]
+        )
+    ops.extend([transforms.ToTensor(), transforms.Normalize(mean=mean, std=std)])
+    return transforms.Compose(ops)
+
+
+def build_dataset(
+    config: ExperimentConfig,
+    split_name: Optional[str],
+    transform,
+    return_metadata: bool = False,
+) -> tuple[CSVImageDataset, list[str]]:
+    df = read_metadata(config.csv_path)
+    class_names = infer_class_names(df, config.label_column, config.class_names)
+    class_to_idx, class_names = class_mapping(class_names)
+    rows = filter_split(df, config.split_column, split_name)
+    dataset = CSVImageDataset(
+        rows=rows,
+        image_dir=config.image_dir,
+        filename_column=config.filename_column,
+        label_column=config.label_column,
+        class_to_idx=class_to_idx,
+        transform=transform,
+        return_metadata=return_metadata,
+    )
+    return dataset, class_names
+
+
+def build_dataloaders(config: ExperimentConfig) -> tuple[DataLoader, DataLoader, list[str]]:
+    train_transform = build_transforms(
+        image_size=config.image_size,
+        train=True,
+        mean=config.normalize_mean,
+        std=config.normalize_std,
+    )
+    eval_transform = build_transforms(
+        image_size=config.image_size,
+        train=False,
+        mean=config.normalize_mean,
+        std=config.normalize_std,
+    )
+    train_dataset, class_names = build_dataset(config, config.train_split, train_transform)
+    val_dataset, _ = build_dataset(config, config.val_split, eval_transform)
+
+    generator = torch.Generator().manual_seed(config.seed)
+    num_workers = config.num_workers if config.num_workers is not None else 0
+    worker_init = partial(seed_worker, base_seed=config.seed)
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=config.batch_size,
+        shuffle=True,
+        num_workers=num_workers,
+        pin_memory=torch.cuda.is_available(),
+        worker_init_fn=worker_init,
+        generator=generator,
+    )
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=config.batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=torch.cuda.is_available(),
+        worker_init_fn=worker_init,
+    )
+    return train_loader, val_loader, class_names

visionkit/evaluate.py

@@ -0,0 +1,111 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional
+
+import pandas as pd
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
+
+from visionkit.data import build_dataset, build_transforms
+from visionkit.metrics import classification_metrics
+
+
+def _record_at(records, index: int) -> dict:
+    if not isinstance(records, dict):
+        return dict(records[index])
+    row = {}
+    for key, value in records.items():
+        item = value[index]
+        if hasattr(item, "item"):
+            try:
+                item = item.item()
+            except ValueError:
+                pass
+        row[key] = item
+    return row
+
+
+def predict(
+    model: torch.nn.Module,
+    dataloader: DataLoader,
+    class_names: list[str],
+    device: Optional[str] = None,
+    positive_index: Optional[int] = None,
+) -> pd.DataFrame:
+    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+    positive_index = positive_index if positive_index is not None else min(1, len(class_names) - 1)
+    model.to(device)
+    model.eval()
+    rows = []
+
+    with torch.no_grad():
+        for images, labels, records in tqdm(dataloader, desc="predict"):
+            images = images.to(device)
+            outputs = model(images)
+            probabilities = F.softmax(outputs, dim=1)
+            predicted = probabilities.argmax(dim=1).cpu().tolist()
+            scores = probabilities[:, positive_index].cpu().tolist()
+
+            for idx in range(images.size(0)):
+                record = _record_at(records, idx)
+                true_label_idx = labels[idx].item() if labels[idx] is not None else None
+                row = dict(record)
+                row.update(
+                    {
+                        "true_label_num": true_label_idx,
+                        "true_label": class_names[true_label_idx] if true_label_idx is not None else None,
+                        "predicted_class_num": predicted[idx],
+                        "predicted_class": class_names[predicted[idx]],
+                        "positive_class_probability": scores[idx],
+                    }
+                )
+                for class_idx, class_name in enumerate(class_names):
+                    row[f"probability_{class_name}"] = probabilities[idx, class_idx].item()
+                rows.append(row)
+    return pd.DataFrame(rows)
+
+
+def evaluate_dataframe(
+    predictions: pd.DataFrame,
+    positive_index: int = 1,
+    n_bootstraps: int = 1000,
+    seed: int = 123,
+) -> dict[str, float]:
+    if predictions.empty or predictions["true_label_num"].nunique() < 2:
+        return {}
+    return classification_metrics(
+        predictions["true_label_num"].astype(int),
+        predictions["predicted_class_num"].astype(int),
+        predictions["positive_class_probability"],
+        positive_index=positive_index,
+        n_bootstraps=n_bootstraps,
+        seed=seed,
+    )
+
+
+def evaluate_model(
+    model: torch.nn.Module,
+    config,
+    split_name: Optional[str],
+    class_names: list[str],
+    output_csv: Path,
+    positive_index: int = 1,
+) -> tuple[pd.DataFrame, dict[str, float]]:
+    transform = build_transforms(
+        image_size=config.image_size,
+        train=False,
+        mean=config.normalize_mean,
+        std=config.normalize_std,
+    )
+    dataset, _ = build_dataset(config, split_name, transform, return_metadata=True)
+    loader = DataLoader(dataset, batch_size=config.batch_size, shuffle=False, num_workers=0)
+    predictions = predict(model, loader, class_names, device=config.device, positive_index=positive_index)
+    output_csv.parent.mkdir(parents=True, exist_ok=True)
+    predictions.to_csv(output_csv, index=False)
+    metrics = evaluate_dataframe(predictions, positive_index=positive_index, seed=config.seed)
+    if metrics:
+        pd.DataFrame([metrics]).to_csv(output_csv.with_name(output_csv.stem + "_metrics.csv"), index=False)
+    return predictions, metrics

visionkit/gradcam.py

@@ -0,0 +1,193 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional, Sequence, Union
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from tqdm.auto import tqdm
+
+from visionkit.data import build_transforms, class_mapping, filter_split, infer_class_names, read_metadata
+from visionkit.evaluate import evaluate_dataframe
+from visionkit.models import get_module_by_path
+
+
+class GradCAM:
+    def __init__(self, model: torch.nn.Module, target_layer: torch.nn.Module) -> None:
+        self.model = model
+        self.target_layer = target_layer
+        self.gradients = None
+        self.activations = None
+        self.handles = []
+        self._register_hooks()
+
+    def _register_hooks(self) -> None:
+        def forward_hook(module, inputs, output):
+            self.activations = output.detach()
+
+        def backward_hook(module, grad_input, grad_output):
+            self.gradients = grad_output[0].detach()
+
+        self.handles.append(self.target_layer.register_forward_hook(forward_hook))
+        self.handles.append(self.target_layer.register_full_backward_hook(backward_hook))
+
+    def generate(self, input_tensor: torch.Tensor, target_class: Optional[int] = None) -> np.ndarray:
+        output = self.model(input_tensor)
+        if target_class is None:
+            target_class = int(output.argmax(dim=1).item())
+        self.model.zero_grad()
+        output[0, target_class].backward()
+
+        if self.gradients is None or self.activations is None:
+            raise RuntimeError("Grad-CAM hooks did not capture gradients or activations.")
+
+        gradients = self.gradients[0]
+        activations = self.activations[0]
+        weights = gradients.mean(dim=(1, 2), keepdim=True)
+        cam = torch.sum(weights * activations, dim=0)
+        cam = F.relu(cam)
+        cam = F.interpolate(
+            cam.unsqueeze(0).unsqueeze(0),
+            size=input_tensor.shape[2:],
+            mode="bilinear",
+            align_corners=False,
+        )
+        cam_np = cam.squeeze().detach().cpu().numpy()
+        denominator = cam_np.max() - cam_np.min()
+        if denominator <= 1e-12:
+            return np.zeros_like(cam_np)
+        return (cam_np - cam_np.min()) / denominator
+
+    def close(self) -> None:
+        for handle in self.handles:
+            handle.remove()
+        self.handles.clear()
+
+
+def _denormalize(image_tensor: torch.Tensor, mean: Sequence[float], std: Sequence[float]) -> np.ndarray:
+    image = image_tensor.detach().cpu().clone()
+    for channel, (mean_value, std_value) in enumerate(zip(mean, std)):
+        image[channel] = image[channel] * std_value + mean_value
+    image = image.clamp(0, 1).permute(1, 2, 0).numpy()
+    return image
+
+
+def save_gradcam_figure(
+    original_image: np.ndarray,
+    cam: np.ndarray,
+    save_path: Path,
+    title: str,
+    alpha: float = 0.4,
+) -> None:
+    save_path.parent.mkdir(parents=True, exist_ok=True)
+    fig, axes = plt.subplots(1, 3, figsize=(12, 4))
+    axes[0].imshow(original_image)
+    axes[0].set_title(title)
+    axes[1].imshow(cam, cmap="jet")
+    axes[1].set_title("Grad-CAM")
+    axes[2].imshow(original_image)
+    axes[2].imshow(cam, cmap="jet", alpha=alpha)
+    axes[2].set_title("Overlay")
+    for axis in axes:
+        axis.axis("off")
+    fig.tight_layout()
+    fig.savefig(save_path, dpi=150)
+    plt.close(fig)
+
+
+def generate_gradcam_report(
+    model: torch.nn.Module,
+    csv_path: Path,
+    image_dir: Path,
+    output_dir: Path,
+    output_csv: Path,
+    filename_column: str = "filename",
+    label_column: str = "label",
+    split_column: Optional[str] = "split",
+    split_name: Optional[str] = None,
+    class_names: Optional[Sequence[str]] = None,
+    target_class: Optional[Union[str, int]] = None,
+    target_layer: Optional[str] = None,
+    image_size: tuple[int, int] = (300, 300),
+    mean: tuple[float, float, float] = (0.485, 0.456, 0.406),
+    std: tuple[float, float, float] = (0.229, 0.224, 0.225),
+    device: Optional[str] = None,
+) -> tuple[pd.DataFrame, dict[str, float]]:
+    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    model.eval()
+    df = filter_split(read_metadata(csv_path), split_column, split_name)
+    class_names = infer_class_names(df, label_column, class_names)
+    class_to_idx, class_names = class_mapping(class_names)
+    positive_index = min(1, len(class_names) - 1)
+    if target_class is None:
+        target_index = positive_index
+    elif isinstance(target_class, int):
+        target_index = target_class
+    else:
+        target_index = class_to_idx[str(target_class)]
+
+    transform = build_transforms(image_size=image_size, train=False, mean=mean, std=std)
+    target_module = get_module_by_path(model, target_layer)
+    gradcam = GradCAM(model, target_module)
+    output_dir = Path(output_dir)
+    output_csv = Path(output_csv)
+    results = []
+
+    try:
+        for record in tqdm(df.to_dict(orient="records"), desc="gradcam"):
+            filename = str(record[filename_column])
+            image_path = Path(image_dir) / filename
+            if not image_path.exists():
+                continue
+            image = Image.open(image_path).convert("RGB")
+            image_tensor = transform(image).unsqueeze(0).to(device)
+            cam = gradcam.generate(image_tensor, target_class=target_index)
+
+            with torch.no_grad():
+                outputs = model(image_tensor)
+                probabilities = F.softmax(outputs, dim=1)[0].detach().cpu()
+                predicted_index = int(probabilities.argmax().item())
+
+            true_label = str(record[label_column]) if label_column in record else None
+            true_index = class_to_idx[true_label] if true_label in class_to_idx else None
+            original = _denormalize(image_tensor[0], mean, std)
+            pred_label = class_names[predicted_index]
+            save_name = f"{Path(filename).stem}__pred-{pred_label}.png"
+            save_gradcam_figure(
+                original,
+                cam,
+                output_dir / save_name,
+                title=f"{filename}\ntrue={true_label} pred={pred_label}",
+            )
+
+            result = dict(record)
+            result.update(
+                {
+                    "gradcam_file": save_name,
+                    "target_class_num": target_index,
+                    "target_class": class_names[target_index],
+                    "true_label_num": true_index,
+                    "true_label": true_label,
+                    "predicted_class_num": predicted_index,
+                    "predicted_class": pred_label,
+                    "positive_class_probability": float(probabilities[positive_index].item()),
+                }
+            )
+            for class_idx, class_name in enumerate(class_names):
+                result[f"probability_{class_name}"] = float(probabilities[class_idx].item())
+            results.append(result)
+    finally:
+        gradcam.close()
+
+    results_df = pd.DataFrame(results)
+    output_csv.parent.mkdir(parents=True, exist_ok=True)
+    results_df.to_csv(output_csv, index=False)
+    metrics = evaluate_dataframe(results_df, positive_index=positive_index)
+    if metrics:
+        pd.DataFrame([metrics]).to_csv(output_csv.with_name(output_csv.stem + "_metrics.csv"), index=False)
+    return results_df, metrics