rgb-to-segmentation 0.1.6__py3-none-any.whl

rgb_to_segmentation/models/base_classifier.py

@@ -0,0 +1,77 @@
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+
+
+class PixelClassifier(pl.LightningModule):
+    """
+    Base class for pixel-wise classifiers.
+    Provides common interface for different model architectures.
+    """
+
+    def __init__(self, output_dim):
+        super().__init__()
+        self.output_dim = output_dim
+        self.loss_fn = nn.CrossEntropyLoss()
+
+    def forward(self, x: torch.Tensor):
+        """Forward pass - to be implemented by subclasses."""
+        raise NotImplementedError
+
+    def image_to_batch(self, x: torch.Tensor):
+        """Convert image tensor to batch for processing.
+
+        Args:
+            x: Input tensor of shape (C, H, W) or (B, C, H, W)
+
+        Returns:
+            Batch tensor ready for model input
+        """
+        raise NotImplementedError
+
+    def _align_logits_and_target(
+        self, logits: torch.Tensor, target: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Reshape predictions/targets so CrossEntropyLoss receives the expected shapes."""
+
+        # Handle channel-first 4D logits from CNN decoders: (B, C, H, W)
+        if logits.dim() == 4:
+            # Targets may come in as (B, 1, H, W) – squeeze the class channel
+            if target.dim() == 4 and target.size(1) == 1:
+                target = target.squeeze(1)
+            return logits, target
+
+        # Everything else (e.g. flattened pixel classifiers): collapse batch/pixel dims
+        logits = logits.view(-1, logits.size(-1))
+        target = target.view(-1)
+        return logits, target
+
+    def training_step(self, batch, batch_idx):
+        """Training step."""
+        sample, target = batch
+        # Apply model-specific batching on GPU
+        sample = self.image_to_batch(sample)
+        target = self.image_to_batch(target).long()
+        
+        logits = self(sample)
+        logits, target = self._align_logits_and_target(logits, target)
+        loss = self.loss_fn(logits, target)
+        self.log("train_loss", loss)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        """Validation step."""
+        sample, target = batch
+        # Apply model-specific batching on GPU
+        sample = self.image_to_batch(sample)
+        target = self.image_to_batch(target).long()
+        
+        logits = self(sample)
+        logits, target = self._align_logits_and_target(logits, target)
+        loss = self.loss_fn(logits, target)
+        self.log("val_loss", loss, prog_bar=True)
+        return loss
+
+    def configure_optimizers(self):
+        """Configure optimizer."""
+        return torch.optim.Adam(self.parameters())

rgb_to_segmentation/models/cnn_decoder.py

@@ -0,0 +1,103 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .base_classifier import PixelClassifier
+
+
+class CNNDecoder(PixelClassifier):
+    """
+    CNN-based pixel classifier for image segmentation.
+
+    Uses a convolutional encoder-decoder architecture to process images
+    and produce per-pixel classifications.
+    """
+
+    def __init__(self, input_channels=3, hidden_dim=64, output_dim=2):
+        """
+        Initialize CNN decoder.
+
+        Args:
+            input_channels: Number of input channels (e.g., 3 for RGB)
+            hidden_dim: Number of hidden channels in conv layers
+            output_dim: Number of output classes
+        """
+        super().__init__(output_dim=output_dim)
+        self.save_hyperparameters()
+
+        self.input_channels = input_channels
+        self.hidden_dim = hidden_dim
+
+        # Encoder: downsampling with convolutions
+        self.encoder = nn.Sequential(
+            nn.Conv2d(input_channels, hidden_dim, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(hidden_dim, hidden_dim * 2, kernel_size=3, stride=2, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(
+                hidden_dim * 2, hidden_dim * 4, kernel_size=3, stride=2, padding=1
+            ),
+            nn.ReLU(),
+        )
+
+        # Decoder: upsampling with transpose convolutions
+        self.decoder = nn.Sequential(
+            nn.ConvTranspose2d(
+                hidden_dim * 4, hidden_dim * 2, kernel_size=4, stride=2, padding=1
+            ),
+            nn.ReLU(),
+            nn.ConvTranspose2d(
+                hidden_dim * 2, hidden_dim, kernel_size=4, stride=2, padding=1
+            ),
+            nn.ReLU(),
+            nn.Conv2d(hidden_dim, output_dim, kernel_size=3, padding=1),
+        )
+
+    def forward(self, x: torch.Tensor):
+        """
+        Forward pass.
+
+        Args:
+            x: Input tensor of shape (B, C, H, W)
+
+        Returns:
+            Logits of shape (B, num_classes, H, W)
+        """
+        encoded = self.encoder(x)
+        decoded = self.decoder(encoded)
+        return decoded
+
+    def image_to_batch(self, x: torch.Tensor):
+        """
+        Convert image tensor to batch format for CNN processing.
+
+        Args:
+            x: Input tensor of shape (C, H, W)
+
+        Returns:
+            Batch tensor of shape (1, C, H, W)
+        """
+        if len(x.shape) == 3:
+            x = x.unsqueeze(0)
+        return x.to(self.device)
+
+    def training_step(self, batch, batch_idx):
+        """Training step with softmax cross entropy loss."""
+        sample, target = batch
+        # sample shape: (B, C, H, W)
+        # target shape: (B, H, W) with class indices
+
+        logits = self(sample)  # (B, num_classes, H, W)
+
+        # Compute loss
+        loss = self.loss_fn(logits, target)
+        self.log("train_loss", loss)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        """Validation step with softmax cross entropy loss."""
+        sample, target = batch
+        logits = self(sample)  # (B, num_classes, H, W)
+        loss = self.loss_fn(logits, target)
+        self.log("val_loss", loss, prog_bar=True)
+        return loss

rgb_to_segmentation/models/pixelwise_classifier.py

@@ -0,0 +1,31 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .base_classifier import PixelClassifier
+
+
+class PixelwiseClassifier(PixelClassifier):
+    step = 0
+    val_step = 0
+
+    def __init__(self, input_dim, hidden_dim, output_dim):
+        super().__init__(output_dim=output_dim)
+        self.save_hyperparameters()
+
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            nn.ReLU(),
+            nn.Linear(hidden_dim, output_dim),
+        )
+
+    def forward(self, x: torch.Tensor):
+        return self.net(x)
+
+    def image_to_batch(self, x: torch.Tensor):
+        if len(x.shape) == 3:
+            x = x.unsqueeze(0)
+        return x.permute(0, 2, 3, 1).reshape(-1, x.shape[1]).to(self.device)
+
+    def batch_to_image(self, batch: torch.Tensor, height: int, width: int):
+        return batch.reshape(1, height, width, 3).permute(0, 3, 1, 2)

rgb_to_segmentation/nn.py

@@ -0,0 +1,164 @@
+import os
+
+import numpy as np
+import torch
+
+from PIL import Image
+from torchvision.io import read_image
+
+from .models.pixelwise_classifier import PixelwiseClassifier
+from .models.cnn_decoder import CNNDecoder
+
+
+def clean_image_nn(
+    image_array: np.ndarray,
+    model: object,
+    colour_map: dict,
+    output_type: str = "rgb",
+) -> np.ndarray:
+    """
+    Clean a single image using the neural network model.
+
+    Args:
+        image_array: Numpy array (H, W, 3) uint8
+        model: Trained model with `image_to_batch` and `forward`
+        colour_map: Dict[int, (r,g,b)] used for RGB output mapping
+        output_type: 'rgb' for colour image, 'index' for class-index mask
+
+    Returns:
+        np.ndarray: (H,W,3) uint8 if output_type='rgb'; else (H,W) uint8
+    """
+    if output_type not in ("rgb", "index"):
+        raise ValueError("output_type must be 'rgb' or 'index'")
+
+    import torch
+
+    img_t = torch.from_numpy(image_array).permute(2, 0, 1).float() / 127.5 - 1.0
+    h, w = img_t.shape[1], img_t.shape[2]
+    batch = model.image_to_batch(img_t)
+
+    with torch.no_grad():
+        probs = model(batch)
+        predicted = torch.argmax(probs, dim=-1).reshape(h, w)
+
+    if output_type == "rgb":
+        rgb_image_t = map_int_to_rgb(predicted, colour_map)
+        return rgb_image_t.permute(1, 2, 0).numpy()
+    else:
+        return predicted.cpu().numpy().astype(np.uint8)
+
+
+def load_model(model_path: str):
+    if os.path.isdir(model_path):
+        checkpoint_files = [f for f in os.listdir(model_path) if f.endswith(".ckpt")]
+
+        if len(checkpoint_files) == 0:
+            raise ValueError(f"No checkpoint files found in directory: {model_path}")
+
+        model_path = os.path.join(model_path, checkpoint_files[0])
+
+    if "pixel_decoder" in model_path:
+        model = PixelwiseClassifier.load_from_checkpoint(checkpoint_path=model_path)
+    elif "cnn_decoder" in model_path:
+        model = CNNDecoder.load_from_checkpoint(checkpoint_path=model_path)
+    else:
+        raise ValueError(
+            "Model path must contain 'pixel_decoder' or 'cnn_decoder' to identify model type."
+        )
+
+    model.eval()
+    return model
+
+
+def map_int_to_rgb(indexed_image: torch.Tensor, colour_map: dict):
+    h, w = indexed_image.shape
+    rgb_image = torch.zeros((3, h, w), dtype=torch.uint8)
+
+    for idx, rgb in colour_map.items():
+        mask = indexed_image == idx
+        for c in range(3):
+            rgb_image[c][mask] = rgb[c]
+
+    return rgb_image
+
+
+def run_inference(
+    input_dir: str,
+    output_dir: str,
+    inplace: bool = False,
+    model_path: str = None,
+    colour_map: dict = None,
+    exts: str = ".png,.jpg,.jpeg,.tiff,.bmp,.gif",
+    name_filter: str = "",
+    output_type: str = "rgb",
+):
+    """
+    Run neural network inference on segmentation images.
+
+    Args:
+        input_dir (str): Path to input directory containing images.
+        output_dir (str, optional): Directory where output images will be written.
+        inplace (bool): Overwrite input images in place.
+        model_path (str): Path to the trained model file.
+        colour_map (dict): Mapping from class indices to RGB tuples.
+        exts (str): Comma-separated list of allowed image extensions.
+        name_filter (str): Only process files whose name contains this substring.
+    """
+    if not inplace and output_dir is None:
+        raise ValueError("Either output_dir must be provided or inplace must be True")
+
+    if model_path is None:
+        raise ValueError("model_path must be provided")
+
+    if colour_map is None:
+        raise ValueError("colour_map must be provided")
+
+    model = load_model(model_path)
+
+    exts_list = [e.lower().strip() for e in exts.split(",")]
+    out_dir = output_dir if not inplace else input_dir
+
+    if not inplace:
+        os.makedirs(out_dir, exist_ok=True)
+
+    print(
+        f"Running inference: input={input_dir} -> output={out_dir}, output_type={output_type}"
+    )
+
+    for root, dirs, files in os.walk(input_dir):
+        rel = os.path.relpath(root, input_dir)
+        out_root = os.path.join(out_dir, rel) if not inplace else root
+
+        os.makedirs(out_root, exist_ok=True)
+
+        for fname in files:
+            if not any(fname.lower().endswith(e) for e in exts_list):
+                continue
+            if name_filter and name_filter not in fname:
+                continue
+
+            in_path = os.path.join(root, fname)
+            out_path = os.path.join(out_root, fname)
+
+            try:
+                # Load as numpy for the shared single-image cleaner
+                pil_img = Image.open(in_path).convert("RGB")
+                np_img = np.array(pil_img, dtype=np.uint8)
+
+                cleaned = clean_image_nn(
+                    np_img, model=model, colour_map=colour_map, output_type=output_type
+                )
+
+                if output_type == "rgb":
+                    pil_image = Image.fromarray(cleaned)
+                    pil_image.save(out_path)
+                elif output_type == "index":
+                    pil_image = Image.fromarray(cleaned.astype("uint8"), mode="L")
+                    pil_image.save(out_path)
+                else:
+                    raise ValueError("output_type must be 'rgb' or 'index'")
+
+            except Exception as e:
+                print(f"Skipping {in_path}: {e}")
+
+    print("Inference done.")

rgb_to_segmentation/train.py

@@ -0,0 +1,212 @@
+import os
+
+import torch
+
+from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
+from torch.utils.data import DataLoader, Dataset, random_split
+from torchvision.io import read_image
+
+from .nn import PixelwiseClassifier
+
+
+VALIDATION_FRACTION = 0.2
+
+
+def map_colour_to_int(sample, colour_map):
+    """Vectorized color to class index mapping."""
+    _, H, W = sample.shape
+
+    # Create lookup tensor: shape (num_classes, 3)
+    num_classes = len(colour_map)
+    color_array = torch.zeros((num_classes, 3), dtype=sample.dtype)
+    for idx, rgb in colour_map.items():
+        color_array[idx] = torch.tensor(rgb, dtype=sample.dtype)
+
+    # Reshape for broadcasting: (3, H, W) -> (H, W, 3) -> (H*W, 3)
+    sample_flat = sample.permute(1, 2, 0).reshape(-1, 3)
+
+    # Compute distances to all colors: (H*W, num_classes)
+    distances = torch.cdist(sample_flat.float(), color_array.float(), p=2)
+
+    # Assign to nearest color
+    image_array = distances.argmin(dim=1).reshape(1, H, W)
+
+    return image_array
+
+
+def collate_fn(batch):
+    """Custom collate to apply image_to_batch on GPU."""
+    samples, targets = zip(*batch)
+    # Stack into batch tensors
+    samples = torch.stack(samples, dim=0)
+    targets = torch.stack(targets, dim=0)
+    # Apply model-specific batching (will happen on GPU in training loop)
+    return samples, targets
+
+
+class SegMaskDataset(Dataset):
+    def __init__(self, paired_filenames, colour_map, model):
+        self.paired_filenames = paired_filenames
+        self.colour_map = colour_map
+        self.to_batch = model.image_to_batch
+
+    def __len__(self):
+        return len(self.paired_filenames)
+
+    def __getitem__(self, index):
+        sample_path = self.paired_filenames[index]["sample"]
+        target_path = self.paired_filenames[index]["target"]
+
+        sample = (
+            read_image(sample_path, mode="RGB") / 127.5
+        ) - 1.0  # Normalize to [-1, 1]
+        target = read_image(target_path, mode="RGB")
+        target = map_colour_to_int(target, self.colour_map)
+        sample = self.to_batch(sample)
+        target = self.to_batch(target).to(torch.long)
+
+        return (sample, target)
+
+
+def get_png_basenames(directory: str) -> list[str]:
+    return [
+        os.path.splitext(filename)[0]
+        for filename in os.listdir(directory)
+        if filename.endswith(".png")
+    ]
+
+
+def get_paired_filenames(
+    image_dir: str,
+    label_dir: str,
+    noisy_basenames: list[str],
+    label_basenames: list[str],
+    training_label_basenames: list[str],
+) -> list[dict]:
+    training_paired_filenames = []
+    val_paired_filenames = []
+
+    for noisy_image_name in noisy_basenames:
+        targets = [
+            target_file
+            for target_file in label_basenames
+            if target_file in noisy_image_name
+        ]
+
+        if len(targets) > 1:
+            raise Exception(
+                f"Multiple target files exist for noisy file {noisy_image_name}: {targets}."
+            )
+        elif len(targets) < 1:
+            print(
+                f"WARNING: No target found for noisy file {noisy_image_name}. Discarding."
+            )
+        else:
+            target_filename = targets[0]
+
+            if target_filename in training_label_basenames:
+                training_paired_filenames.append(
+                    {
+                        "sample": f"{image_dir}/{noisy_image_name}.png",
+                        "target": f"{label_dir}/{target_filename}.png",
+                    }
+                )
+            else:
+                val_paired_filenames.append(
+                    {
+                        "sample": f"{image_dir}/{noisy_image_name}.png",
+                        "target": f"{label_dir}/{target_filename}.png",
+                    }
+                )
+
+    return training_paired_filenames, val_paired_filenames
+
+
+def get_dataloaders(
+    image_dir, label_dir, colour_map, model
+) -> tuple[DataLoader, DataLoader]:
+    noisy_basenames = get_png_basenames(image_dir)
+    label_basenames = get_png_basenames(label_dir)
+
+    training_label_basenames, val_label_basenames = random_split(
+        label_basenames, [1 - VALIDATION_FRACTION, VALIDATION_FRACTION]
+    )
+
+    training_paired_filenames, val_paired_filenames = get_paired_filenames(
+        image_dir,
+        label_dir,
+        noisy_basenames,
+        label_basenames,
+        training_label_basenames,
+    )
+
+    train_dataset = SegMaskDataset(training_paired_filenames, colour_map, model)
+    val_dataset = SegMaskDataset(val_paired_filenames, colour_map, model)
+
+    train_dataloader = DataLoader(
+        dataset=train_dataset,
+        batch_size=1,
+        shuffle=True,
+        num_workers=4,
+        prefetch_factor=4,
+        collate_fn=collate_fn,
+    )
+    val_dataloader = DataLoader(
+        dataset=val_dataset,
+        batch_size=1,
+        shuffle=False,
+        num_workers=4,
+        prefetch_factor=4,
+        collate_fn=collate_fn,
+    )
+
+    return train_dataloader, val_dataloader
+
+
+def get_model(model_type, colour_map):
+    num_classes = len(colour_map.keys())
+
+    if model_type == "pixel_decoder":
+        return PixelwiseClassifier(input_dim=3, hidden_dim=32, output_dim=num_classes)
+
+    raise ValueError(f"Invalid model type: {model_type}")
+
+
+def train_model(
+    image_dir: str,
+    label_dir: str,
+    output_dir: str,
+    colour_map: dict,
+    model_type: str = "pixel_decoder",
+):
+    """
+    Train a neural network model for segmentation cleaning.
+
+    Args:
+        image_dir (str): Path to directory containing noisy images.
+        label_dir (str): Path to directory containing target RGB labels.
+        output_dir (str): Directory where model weights will be saved.
+        colour_map (dict): Mapping from class indices to RGB tuples.
+        model_type (str): The type of model to train.
+    """
+    model = get_model(model_type, colour_map)
+    train_dataloader, val_dataloader = get_dataloaders(
+        image_dir, label_dir, colour_map, model
+    )
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    early_stop = EarlyStopping(monitor="val_loss", mode="min", verbose=True)
+    checkpoint = ModelCheckpoint(
+        monitor="val_loss",
+        mode="min",
+        save_top_k=1,
+        dirpath=output_dir,
+        filename=f"{model_type}.ckpt",
+    )
+    trainer = Trainer(max_epochs=100, callbacks=[early_stop, checkpoint])
+
+    trainer.fit(
+        model, train_dataloaders=train_dataloader, val_dataloaders=val_dataloader
+    )

rgb_to_segmentation/utils.py

@@ -0,0 +1,40 @@
+from typing import List, Tuple
+
+
+def parse_colours_from_string(colours_str: str) -> List[Tuple[int, int, int]]:
+    """Parse a semicolon-separated string of RGB triples into a list of tuples."""
+
+    parts = [p.strip() for p in colours_str.split(";") if p.strip()]
+    colours = []
+
+    for p in parts:
+        rgb = tuple(int(x) for x in p.split(","))
+
+        if len(rgb) != 3:
+            raise ValueError(f"Invalid colour triple: {p}")
+        colours.append(rgb)
+
+    return colours
+
+
+def parse_colours_from_file(path: str) -> List[Tuple[int, int, int]]:
+    """Parse a file with one RGB triple per line into a list of tuples."""
+
+    colours = []
+
+    with open(path, "r") as f:
+        for line in f:
+            line = line.strip()
+            if not line or line.startswith("#"):
+                continue
+
+            rgb = tuple(int(x) for x in line.split(","))
+            if len(rgb) != 3:
+                raise ValueError(f"Invalid colour triple in file {path}: {line}")
+
+            colours.append(rgb)
+
+    if not colours:
+        raise ValueError(f"No colours found in file: {path}")
+
+    return colours