findcrack 0.2.0__py3-none-any.whl

findcrack/__init__.py

@@ -0,0 +1,17 @@
+from .pipeline import CrackInferencePipeline
+from .models import load_model, UNet, DeepCrack, list_models, register_model
+from .metrics import calculate_metrics
+from .preprocess import apply_lab_clahe
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "CrackInferencePipeline",
+    "load_model",
+    "UNet",
+    "DeepCrack",
+    "calculate_metrics",
+    "apply_lab_clahe",
+    "list_models",
+    "register_model",
+]

findcrack/metrics.py

@@ -0,0 +1,23 @@
+import numpy as np
+
+def calculate_metrics(y_true: np.ndarray, y_prediction: np.ndarray, epsilon: float = 1e-7) -> dict:
+    """
+    Calculates IoU, Dice, Precision, Recall, and Pixel Accuracy.
+    """
+    
+    y_true = y_true.astype(bool)
+    y_prediction = y_prediction.astype(bool)
+    
+    # True Positives, False Positives, True Negatives, False Negatives
+    TP = np.sum(y_true & y_prediction)
+    FP = np.sum(~y_true & y_prediction)
+    FN = np.sum(y_true & ~y_prediction)
+    TN = np.sum(~y_true & ~y_prediction)
+    
+    return {
+        "IoU": TP / (TP + FP + FN + epsilon),
+        "Dice": (2 * TP) / (2 * TP + FP + FN + epsilon),
+        "Precision": TP / (TP + FP + epsilon),
+        "Recall": TP / (TP + FN + epsilon),
+        "Pixel Accuracy": (TP + TN) / (TP + TN + FP + FN + epsilon)
+    }

findcrack/models/__init__.py

@@ -0,0 +1,16 @@
+from .unet import UNet
+from .deepcrack import DeepCrack
+from .onnx_wrapper import ONNXModelWrapper
+from .zoo import load_model, MODEL_REGISTRY, list_models, register_model
+
+__all__ = [
+    "UNet",
+    "DeepCrack",
+    "ONNXModelWrapper",
+    "load_model",
+    "MODEL_REGISTRY",
+    "list_models",
+    "register_model",
+]
+
+

findcrack/models/deepcrack.py

@@ -0,0 +1,93 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class DoubleConv(nn.Module):
+    """(convolution => [BN] => ReLU) * 2"""
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.double_conv(x)
+
+
+class DeepCrack(nn.Module):
+    """
+    DeepCrack: A Deep Hierarchical Feature Learning Architecture for Crack Segmentation.
+    Fuses hierarchical convolutional features from both the encoder and decoder stages 
+    at the same scale.
+    """
+    def __init__(self, n_channels: int = 3, n_classes: int = 1):
+        super().__init__()
+        self.n_channels = n_channels
+        self.n_classes = n_classes
+
+        # Encoder (downsampling blocks)
+        self.enc1 = DoubleConv(n_channels, 64)
+        self.enc2 = DoubleConv(64, 128)
+        self.enc3 = DoubleConv(128, 256)
+        self.enc4 = DoubleConv(256, 512)
+        self.enc5 = DoubleConv(512, 512)
+
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+
+        # Decoder (upsampling & concatenation blocks)
+        self.dec5 = DoubleConv(512, 512)
+        self.dec4 = DoubleConv(512 + 512, 256)
+        self.dec3 = DoubleConv(256 + 256, 128)
+        self.dec2 = DoubleConv(128 + 128, 64)
+        self.dec1 = DoubleConv(64 + 64, 64)
+
+        # Side prediction layers (maps feature channels to class channels at each scale)
+        self.side1 = nn.Conv2d(64, n_classes, kernel_size=1)
+        self.side2 = nn.Conv2d(64, n_classes, kernel_size=1)
+        self.side3 = nn.Conv2d(128, n_classes, kernel_size=1)
+        self.side4 = nn.Conv2d(256, n_classes, kernel_size=1)
+        self.side5 = nn.Conv2d(512, n_classes, kernel_size=1)
+
+        # Fusion layer that combines all 5 side predictions into the final output
+        self.fuse = nn.Conv2d(n_classes * 5, n_classes, kernel_size=1)
+
+    def forward(self, x):
+        # 1. Encoder path
+        e1 = self.enc1(x)
+        e2 = self.enc2(self.pool(e1))
+        e3 = self.enc3(self.pool(e2))
+        e4 = self.enc4(self.pool(e3))
+        e5 = self.enc5(self.pool(e4))
+
+        # 2. Decoder path (with bilinear interpolation upsampling)
+        d5 = self.dec5(e5)
+        
+        d4_up = F.interpolate(d5, size=e4.shape[2:], mode='bilinear', align_corners=True)
+        d4 = self.dec4(torch.cat([d4_up, e4], dim=1))
+
+        d3_up = F.interpolate(d4, size=e3.shape[2:], mode='bilinear', align_corners=True)
+        d3 = self.dec3(torch.cat([d3_up, e3], dim=1))
+
+        d2_up = F.interpolate(d3, size=e2.shape[2:], mode='bilinear', align_corners=True)
+        d2 = self.dec2(torch.cat([d2_up, e2], dim=1))
+
+        d1_up = F.interpolate(d2, size=e1.shape[2:], mode='bilinear', align_corners=True)
+        d1 = self.dec1(torch.cat([d1_up, e1], dim=1))
+
+        # 3. Extract side predictions and upsample to input dimensions
+        h, w = x.shape[2:]
+        s1 = F.interpolate(self.side1(d1), size=(h, w), mode='bilinear', align_corners=True)
+        s2 = F.interpolate(self.side2(d2), size=(h, w), mode='bilinear', align_corners=True)
+        s3 = F.interpolate(self.side3(d3), size=(h, w), mode='bilinear', align_corners=True)
+        s4 = F.interpolate(self.side4(d4), size=(h, w), mode='bilinear', align_corners=True)
+        s5 = F.interpolate(self.side5(d5), size=(h, w), mode='bilinear', align_corners=True)
+
+        # 4. Fuse side predictions
+        fused = self.fuse(torch.cat([s1, s2, s3, s4, s5], dim=1))
+
+        return fused

findcrack/models/onnx_wrapper.py

@@ -0,0 +1,35 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+class ONNXModelWrapper(nn.Module):
+    """
+    Wraps an ONNX Runtime InferenceSession inside a PyTorch nn.Module.
+    This allows running inference on ONNX models using the exact same code
+    and APIs as PyTorch models, supporting both CPU/GPU tensor operations
+    and test-time augmentation (TTA) pipelines.
+    """
+    def __init__(self, model_path: str, device: str = "cpu"):
+        super().__init__()
+        import onnxruntime as ort
+        
+        # Select execution providers based on the target device
+        if device == "cuda" or (isinstance(device, torch.device) and device.type == "cuda"):
+            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+        else:
+            providers = ["CPUExecutionProvider"]
+            
+        self.session = ort.InferenceSession(model_path, providers=providers)
+        self.input_name = self.session.get_inputs()[0].name
+        self.output_name = self.session.get_outputs()[0].name
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # 1. Convert PyTorch tensor to NumPy array (typically expected as float32)
+        x_np = x.detach().cpu().numpy().astype(np.float32)
+        
+        # 2. Run inference using ONNX Runtime
+        outputs = self.session.run([self.output_name], {self.input_name: x_np})
+        
+        # 3. Convert prediction back to PyTorch tensor and move to the original device
+        out_tensor = torch.from_numpy(outputs[0]).to(x.device)
+        return out_tensor

findcrack/models/unet.py

@@ -0,0 +1,103 @@
+import torch
+import torch.nn as nn
+
+# U-Net Model Definition
+class DoubleConv(nn.Module):
+    """ (convolution => [BN] => ReLU) * 2 """
+    def __init__(self, in_channels, out_channels, mid_channels=None):
+        super().__init__()
+        if not mid_channels:
+            mid_channels = out_channels
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(mid_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self, x):
+        return self.double_conv(x)
+
+
+class Down(nn.Module):
+    """ Downscaling with maxpool then double conv """
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        self.maxpool_conv = nn.Sequential(
+            nn.MaxPool2d(2),
+            DoubleConv(in_channels, out_channels)
+        )
+
+    def forward(self, x):
+        return self.maxpool_conv(x)
+
+
+class Up(nn.Module):
+    """ Upscaling then double conv """
+    def __init__(self, in_channels, out_channels, bilinear=True):
+        super().__init__()
+
+        if bilinear:
+            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
+        else:
+            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
+            self.conv = DoubleConv(in_channels, out_channels)
+
+
+    def forward(self, x1, x2):
+        x1 = self.up(x1)
+        # Input tensor might not be perfectly divisible by 2, so crop x2 to match x1's size
+        # This handles potential size mismatches due to padding in earlier layers or odd input dimensions
+        diffY = x2.size()[2] - x1.size()[2]
+        diffX = x2.size()[3] - x1.size()[3]
+
+        # Pad x1 if necessary to match x2's size (or crop x2 if it's larger)
+        x1 = nn.functional.pad(x1, [diffX // 2, diffX - diffX // 2,
+                                    diffY // 2, diffY - diffY // 2])
+        x = torch.cat([x2, x1], dim=1)
+        return self.conv(x)
+
+
+class OutConv(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(OutConv, self).__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class UNet(nn.Module):
+    def __init__(self, n_channels, n_classes, bilinear=False):
+        super(UNet, self).__init__()
+        self.n_channels = n_channels
+        self.n_classes = n_classes
+        self.bilinear = bilinear
+
+        self.inc = DoubleConv(n_channels, 64)
+        self.down1 = Down(64, 128)
+        self.down2 = Down(128, 256)
+        self.down3 = Down(256, 512)
+        factor = 2 if bilinear else 1
+        self.down4 = Down(512, 1024 // factor)
+        self.up1 = Up(1024, 512 // factor, bilinear)
+        self.up2 = Up(512, 256 // factor, bilinear)
+        self.up3 = Up(256, 128 // factor, bilinear)
+        self.up4 = Up(128, 64, bilinear)
+        self.outc = OutConv(64, n_classes)
+
+    def forward(self, x):
+        x1 = self.inc(x)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        x = self.up1(x5, x4) # F.pad needs to be imported or handled in Up()
+        x = self.up2(x, x3)
+        x = self.up3(x, x2)
+        x = self.up4(x, x1)
+        logits = self.outc(x)
+        return logits

findcrack/models/zoo.py

@@ -0,0 +1,192 @@
+import os
+import torch
+import hashlib
+from torch.hub import get_dir, download_url_to_file
+
+from .unet import UNet
+from .deepcrack import DeepCrack
+from .onnx_wrapper import ONNXModelWrapper
+
+# Model registry detailing architectures, default arguments, remote URLs, hashes, and backend type.
+MODEL_REGISTRY = {
+    "Seg_UNET_CFD_actual_v1": {
+        "metadata": {
+            "loss_functions": ["TverskyLoss"],
+            "input_shape": [3, 512, 512],
+        },
+        "architecture": UNet,
+        "kwargs": {
+            "n_channels": 3,
+            "n_classes": 1,
+            "bilinear": False
+        },
+        "backend": "pytorch",
+        "url": "https://github.com/StrikerEurika/findcrack/releases/download/v0.2.0/Seg_UNET_CFD_actual_v1_best.pth",
+        "sha256": None  # Users can supply checksums to verify integrity
+    },
+    "Seg_UNET_CFD_actual_v2": {
+        "metadata": {
+            "loss_functions": ["BCEWithLogitsLoss", "DiceLoss"],
+            "input_shape": [3, 512, 512],
+        },
+        "architecture": UNet,
+        "kwargs": {
+            "n_channels": 3,
+            "n_classes": 1,
+            "bilinear": False
+        },
+        "backend": "pytorch",
+        "url": "https://github.com/StrikerEurika/findcrack/releases/download/v0.2.0/Seg_UNET_CFD_actual_v2_best.pth",
+        "sha256": None  # Users can supply checksums to verify integrity
+    },
+}
+
+def verify_sha256(filepath: str, expected_hash: str) -> bool:
+    """Verifies file integrity via SHA256 checksum."""
+    if not expected_hash:
+        return True
+    sha256_hash = hashlib.sha256()
+    with open(filepath, "rb") as f:
+        for byte_block in iter(lambda: f.read(4096), b""):
+            sha256_hash.update(byte_block)
+    return sha256_hash.hexdigest() == expected_hash
+
+def get_cache_dir() -> str:
+    """Returns the package's weight caching directory."""
+    hub_directory = get_dir()
+    model_directory = os.path.join(hub_directory, 'checkpoints', 'findcrack')
+    os.makedirs(model_directory, exist_ok=True)
+    return model_directory
+
+def download_weight_file(url: str, cached_file: str, expected_hash: str = None):
+    """Downloads remote weight file and verifies its checksum."""
+    print(f"Downloading weights from {url}... (This is a one-time download)")
+    try:
+        download_url_to_file(url, cached_file, progress=True)
+    except Exception as e:
+        raise IOError(f"Failed to download weight file from {url}. Error: {e}")
+        
+    if expected_hash and not verify_sha256(cached_file, expected_hash):
+        if os.path.exists(cached_file):
+            os.remove(cached_file)
+        raise ValueError(f"Checksum verification failed for {cached_file}. Cache cleared.")
+
+def load_model(
+    variant: str,
+    device: str = "cpu",
+    force_download: bool = False,
+    architecture = None,
+    kwargs: dict = None,
+    backend: str = None,
+    sha256: str = None
+) -> torch.nn.Module:
+    """
+    Loads a model by its registry name OR directly from a remote HTTP(S) URL.
+    
+    Args:
+        variant: Registry name (e.g. 'unet_cfd_v1') OR direct remote URL (e.g. 'https://.../model.pth').
+        device: The target device for execution (e.g., 'cpu', 'cuda', 'mps').
+        force_download: If True, deletes cached files and re-downloads weights.
+        architecture: Model architecture class (e.g., UNet) - required if loading PyTorch weights from a URL.
+        kwargs: Keyword arguments for instantiating the model class (used when loading from a URL).
+        backend: Backend target ('pytorch' or 'onnx'). Automatically inferred from URL if loading from a URL.
+        sha256: Optional SHA256 checksum to verify file integrity.
+    """
+    is_url = variant.startswith("http://") or variant.startswith("https://")
+    
+    if is_url:
+        url = variant
+        if backend is None:
+            backend = "onnx" if url.lower().endswith(".onnx") else "pytorch"
+            
+        if backend == "pytorch" and architecture is None:
+            raise ValueError(
+                "You must specify the 'architecture' class (e.g., UNet, DeepCrack) "
+                "when loading a PyTorch model directly from a URL."
+            )
+            
+        config = {
+            "architecture": architecture,
+            "kwargs": kwargs or {},
+            "backend": backend,
+            "url": url,
+            "sha256": sha256
+        }
+    else:
+        if variant not in MODEL_REGISTRY:
+            available_variants = list(MODEL_REGISTRY.keys())
+            raise ValueError(f"Unknown variant '{variant}'. Available: {available_variants}")
+        config = MODEL_REGISTRY[variant]
+        backend = config.get("backend", "pytorch")
+    
+    # 1. Resolve caching paths
+    filename = os.path.basename(config["url"])
+    cached_file = os.path.join(get_cache_dir(), filename)
+    
+    if force_download and os.path.exists(cached_file):
+        try:
+            os.remove(cached_file)
+        except OSError:
+            pass
+        
+    # 2. Download weights if missing
+    if not os.path.exists(cached_file):
+        download_weight_file(config["url"], cached_file, config.get("sha256"))
+        
+    # 3. Instantiate and load weights based on backend type
+    if backend == "pytorch":
+        arch_class = config["architecture"]
+        model = arch_class(**config.get("kwargs", {}))
+        
+        # Load weights
+        state_dict = torch.load(cached_file, map_location=device)
+        model.load_state_dict(state_dict)
+        return model.to(device).eval()
+        
+    elif backend == "onnx":
+        model = ONNXModelWrapper(cached_file, device=device)
+        return model
+        
+    else:
+        raise ValueError(f"Unsupported backend: {backend}")
+
+
+def list_models() -> list:
+    """
+    Returns a list of all available pre-trained model variants in the registry.
+    """
+    return list(MODEL_REGISTRY.keys())
+
+
+def register_model(
+    name: str,
+    url: str,
+    architecture = None,
+    kwargs: dict = None,
+    backend: str = "pytorch",
+    sha256: str = None
+):
+    """
+    Dynamically registers a custom model variant at runtime.
+    This allows users to define custom remote models and use the standard load_model API.
+    
+    Args:
+        name: Name/identifier of the variant.
+        url: Remote URL to download the model weights/file from.
+        architecture: The PyTorch class/type of the model (not required for ONNX).
+        kwargs: Keyword arguments for instantiating the model class.
+        backend: Backend framework, either 'pytorch' or 'onnx'.
+        sha256: Optional SHA256 checksum to verify the downloaded file.
+    """
+    if backend not in ("pytorch", "onnx"):
+        raise ValueError("backend must be 'pytorch' or 'onnx'")
+    if backend == "pytorch" and architecture is None:
+        raise ValueError("architecture class must be provided for PyTorch backend")
+        
+    MODEL_REGISTRY[name] = {
+        "architecture": architecture,
+        "kwargs": kwargs or {},
+        "backend": backend,
+        "url": url,
+        "sha256": sha256
+    }

findcrack/patching.py

@@ -0,0 +1,172 @@
+import numpy as np
+from typing import Tuple, Generator
+
+"""
+Version 1
+"""
+class PatchExtractor:
+    """
+    Extracts the overlapping patches from a large image.
+    """
+    def __init__(self, patch_size: Tuple[int, int], overlap_ratio: float = 0.2):
+        """
+        Args:
+            patch_size (height, width): the size of the patch to be extracted.
+            overlap_ratio: float number between 0.0 and 0.99. the default value is 0.2,
+            meaning that the patches will overlap by 20% in both dimensions. 
+        """
+        
+        if not (0.0 <= overlap_ratio < 1.0):
+            raise ValueError("overlap_ratio must be between 0.0 and 1.0")
+        
+        self.patch_height, self.patch_width = patch_size
+        self.stride_height = int(self.patch_height * (1 - overlap_ratio))
+        self.stride_width = int(self.patch_width * (1 - overlap_ratio))
+        
+        # ensure that stride is at least 1 to avoid infinite loops
+        self.stride_height = max(1, self.stride_height)
+        self.stride_width = max(1, self.stride_width)
+
+    def extract(self, image: np.ndarray) -> Generator[Tuple[np.ndarray, Tuple[int, int]], None, None]:
+        """
+        Yields patches and their top-left (y, x) coordinates.
+        Handles edges by shifting the last patch to align with the image border.
+        """
+        
+        image_height, image_width = image.shape[:2]
+        seen_coordinates = set()
+        
+        for y in range(0, image_height, self.stride_height):
+            for x in range(0, image_width, self.stride_width):
+                
+                # shift the patch if it goes out of bounds
+                if y + self.patch_height > image_height:
+                    y = image_height - self.patch_height
+                
+                if x + self.patch_width > image_width:
+                    x = image_width - self.patch_width
+                    
+                # ensure we don't yield the same patch multiple times
+                if (y, x) in seen_coordinates:
+                    continue
+                seen_coordinates.add((y, x))
+
+                # yield the patch and its coordinates
+                yield image[y:y+self.patch_height, x:x+self.patch_width], (y, x)
+                
+# Patch Reconstruction
+class PatchBlender:
+    """
+    Reconstructs the full image from the overlapping patches using Gaussian Blending
+    to eliminate seam artifacts
+    """
+    
+    def __init__(self, output_shape: Tuple[int, int], patch_size: Tuple[int, int], num_classes: int = 1):
+        self.output_height, self.output_width = output_shape
+        self.patch_height, self.patch_width = patch_size
+        self.num_classes = num_classes
+        
+        # accumulater for value and weights
+        # using float32 to avoid overflow during accumulation
+        self.accumulated_values = np.zeros((num_classes, self.output_height, self.output_width), dtype=np.float32)
+        self.accumilated_weights = np.zeros((num_classes, self.output_height, self.output_width), dtype=np.float32)
+
+        # pre-compute the 2d gaussian wieghts for a single patch
+        self.weight_map = self._create_gaussian_weight(self.patch_height, self.patch_width)
+
+
+    def _create_gaussian_weight(self, height: int, width: int, sigma_factor: float = 0.25) -> np.ndarray:
+        """
+        Create a 2d Gaussian mask. Center center is 1.0, edge fade to 0.0
+        """
+        x = np.linspace(-1, 1, width)
+        y = np.linspace(-1, 1, height)
+
+        X, Y = np.meshgrid(x, y)
+
+        # sigma_factor controls the drop-off. o.25 makes the edges fade to 0.0
+        weights = np.exp(-(X**2 + Y**2) / (2 * sigma_factor**2))
+        return weights.astype(np.float32)
+
+
+    def add_patch(self, patch_prediction: np.ndarray, coordinates: Tuple[int, int]):
+        """
+        Adds a predicte patch to the accomultor.
+        """
+        y, x = coordinates
+        c = patch_prediction.shape[0]
+        
+        # expand the weight map to march number of classes if necessary
+        weights = np.expand_dims(self.weight_map, 0) if c == 1 else np.stack([self.weight_map] * c, axis=0)
+        
+        self.accumulated_values[:, y:y+self.patch_height, x:x+self.patch_width] += patch_prediction * weights
+        self.accumilated_weights[:, y:y+self.patch_height, x:x+self.patch_width] += weights
+        
+    
+    def merge(self) -> np.ndarray:
+        """
+        Return the final blended image of the shape (C, H, W)
+        """
+        
+        # avoid division by zero in areas with no patches
+        self.accumilated_weights[self.accumilated_weights == 0] = 1.0
+        
+        result = self.accumulated_values / self.accumilated_weights
+        
+        # if single class. squeeze the channel for convenience
+        if self.num_classes == 1:
+            return result.squeeze(0)
+        return result
+    
+    
+"""
+Version 2
+"""
+class SlidingWindowExtractor:
+    def __init__(self, patch_size: int, overlap_ratio: float = 0.2):
+        self.patch_size = patch_size
+        self.step_size = max(1, int(patch_size * (1 - overlap_ratio)))
+    
+    def extract(self, image: np.ndarray) -> Generator[Tuple[np.ndarray, Tuple[int, int]], None, None]:
+        """
+        yields patches and their (y, x) coordinates. automatically handles edges.
+        """
+        height, width = image.shape[:2]
+        seen_coordinates = set()
+        
+        for y in range(0, height, self.step_size):
+            for x in range(0, width, self.step_size):
+                # shift the patch if it goes out of bounds
+                if y + self.patch_size > height: y = height - self.patch_size
+                if x + self.patch_size > width: x = width - self.patch_size
+                
+                if (y, x) in seen_coordinates: continue
+                seen_coordinates.add((y, x))
+
+                yield image[y:y+self.patch_size, x:x+self.patch_size], (y, x)
+                
+class CountMapBlender:
+    """
+    Reconstructs the full image by averaging overlapping patches.
+    """
+    def __init__(self, shape: Tuple[int, int]):
+        self.prediction_map = np.zeros(shape, dtype=np.float32)
+        self.count_map = np.zeros(shape, dtype=np.int32)
+        
+    def add(self, patch: np.ndarray, coordinates: Tuple[int, int]):
+        """
+        Adds a patch to the prediction map and updates the count map.
+        """
+        y, x = coordinates
+        height, width = patch.shape
+        
+        self.prediction_map[y:y+height, x:x+width] += patch
+        self.count_map[y:y+height, x:x+width] += 1
+        
+    def merge(self) -> np.ndarray:
+        """
+        Merges the prediction map with the count map to produce the final blended image.
+        """
+        valid = self.count_map > 0
+        self.prediction_map[valid] /= self.count_map[valid]
+        return self.prediction_map

findcrack/pipeline.py

@@ -0,0 +1,89 @@
+import cv2
+import torch
+import numpy as np
+from pathlib import Path
+from PIL import Image
+
+from .preprocess import apply_lab_clahe, get_inference_transform
+from .tta import tta_forward
+from .patching import CountMapBlender, PatchExtractor, PatchBlender, SlidingWindowExtractor
+from .models import load_model
+
+
+class CrackInferencePipeline:
+    """
+    pipeline for running inference on large images.
+    """
+    def __init__(self, model: torch.nn.Module, device: str = "cuda",
+                 patch_size: int = 512, overlap_ratio: float = 0.2, 
+                 confidence_threhold: float = 0.5, use_tta: bool = False):
+        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
+        self.model = model.to(self.device).eval()
+        self.patch_size = patch_size
+        self.overlap_ratio = overlap_ratio
+        self.confidence_threshold = confidence_threhold
+        self.use_tta = use_tta
+        
+        self.extractor = SlidingWindowExtractor(self.patch_size, self.overlap_ratio)
+        self.transform = get_inference_transform()
+        
+    @classmethod
+    def from_checkpoint(cls, model_class, checkpoint_path: str, **kwargs):
+        """
+        Helper function to load model
+        """
+        model = model_class(n_channels=3, n_classes=1)  # Assuming binary segmentation
+        state_dict = torch.load(checkpoint_path, map_location="cpu")
+        model.load_state_dict(state_dict)
+
+        return cls(model, **kwargs)
+    
+    @classmethod
+    def from_pretrained(cls, variant: str, device: str = "cuda", **kwargs):
+        """
+        Helper function to load pretrained model from the model zoo.
+        """
+        model = load_model(variant, device=device)
+        return cls(model, device=device, **kwargs)
+        
+    
+    @torch.no_grad()
+    def predict(self, image_path: str) -> dict:
+        """
+        Runs full inference pipeline on a large image.
+        Returns a dictionary with the original image, probability map, and binary mask.
+        """
+        # Load Image
+        original_image = np.array(Image.open(image_path).convert('RGB'))
+        height, width, _ = original_image.shape
+        
+        # Preprocess (LAB-CLAHE)
+        preprocessed_image = apply_lab_clahe(original_image)
+        
+        # Initialize Blender
+        blender = CountMapBlender(shape=(height, width))
+        
+         # Sliding Window Inference
+        for patch_rgb, coordinates in self.extractor.extract(preprocessed_image):
+            # Transform to tensor
+            transformed = self.transform(image=patch_rgb)
+            patch_tensor = transformed["image"].to(self.device)
+            
+            # Run Model
+            if self.use_tta:
+                pred_prob = tta_forward(self.model, patch_tensor)
+            else:
+                pred_prob = torch.sigmoid(self.model(patch_tensor.unsqueeze(0))).squeeze()
+                
+            # Add to blender
+            blender.add(pred_prob.cpu().numpy(), coordinates)
+            
+        # Merge and Threshold
+        confidence_map = blender.merge()
+        binary_mask = (confidence_map > self.confidence_threshold).astype(np.uint8) * 255
+        
+        return {
+            "original_image": original_image,
+            "confidence_map": confidence_map,
+            "binary_mask": binary_mask
+        }

findcrack/preprocess.py

@@ -0,0 +1,25 @@
+import cv2
+import numpy as np
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+def apply_lab_clahe(image: np.ndarray, clip_limit: float = 2.0) -> np.ndarray:
+    """
+    Apply clahe to the L-channel of the LAB color space to enhance
+    local contrast without altering color balance.
+    """
+    lab_image = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
+    l_channel, a_channel, b_channel = cv2.split(lab_image)
+    clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=(8, 8))
+    cl = clahe.apply(l_channel)
+    limg = cv2.merge((cl, a_channel, b_channel))
+    return cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)
+
+def get_inference_transform():
+    """
+    standard ImageNet normalization required by most pretrained models.
+    """
+    return A.Compose([
+        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+        ToTensorV2(),
+    ])

findcrack/tta.py

@@ -0,0 +1,40 @@
+import torch
+
+def tta_forward(model: torch.nn.Module, image_tensor: torch.Tensor) -> torch.Tensor:
+    """
+    apply 4-way test-time augmentation and averages the sigmoid predictions.
+    expects image tensor to be shape (C, H, W). returns (H, W)
+    """
+    
+    with torch.no_grad():
+        x = image_tensor.unsqueeze(0)
+
+        # original
+        original_prediction = torch.sigmoid(model(x))
+        
+        # horizontal flip
+        horizontal_flip_prediction = torch.flip(torch.sigmoid(model(torch.flip(x, dims=[3]))), dims=[3])
+
+        # vertical flip
+        vertical_flip_prediction = torch.flip(torch.sigmoid(model(torch.flip(x, dims=[2]))), dims=[2])
+
+        # diagonal flip
+        # diagonal_flip_prediction = torch.flip(torch.sigmoid(model(torch.flip(x, dims=[2, 3]))), dims=[2, 3])
+
+        # average the predictions
+        # final_prediction = (original_prediction + horizontal_flip_prediction + vertical_flip_prediction + diagonal_flip_prediction) / 4
+
+        # rotate 90 degrees clockwise
+        rotated_90_prediction = torch.rot90(
+            torch.sigmoid(model(torch.rot90(x, k=1, dims=[2, 3]))),
+            k=-1, dims=[2,3]
+        )
+        
+        # average predictions
+        averaged_prediction = (
+            original_prediction + 
+            horizontal_flip_prediction + 
+            vertical_flip_prediction + 
+            rotated_90_prediction) / 4.0
+
+    return averaged_prediction.squeeze(0)

findcrack-0.2.0.dist-info/METADATA

@@ -0,0 +1,152 @@
+Metadata-Version: 2.4
+Name: findcrack
+Version: 0.2.0
+Summary: A deep learning crack detection package supporting U-Net and DeepCrack models with PyTorch and ONNX backends.
+Author: StrikerEurika
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/StrikerEurika/findcrack
+Project-URL: Repository, https://github.com/StrikerEurika/findcrack
+Project-URL: Releases, https://github.com/StrikerEurika/findcrack/releases
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Image Processing
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+Requires-Dist: albumentations>=2.0.8
+Requires-Dist: onnxruntime>=1.27.0
+Requires-Dist: opencv-python>=4.13.0.92
+Requires-Dist: torch>=2.6.0
+Requires-Dist: torchaudio>=2.6.0
+Requires-Dist: torchvision>=0.21.0
+
+# findcrack
+
+`findcrack` is a deep learning crack detection package designed for pixel-level segmentation on high-resolution images. It supports U-Net and DeepCrack architectures, providing an easy-to-use API for inference, model caching, and multi-backend execution (PyTorch & ONNX).
+
+---
+
+## Features
+
+- **Pre-trained Model Zoo**: Fetch pre-trained model weights (e.g., `Seg_UNET_CFD_actual_v1`, `Seg_UNET_CFD_actual_v2`) dynamically on demand.
+- **Unified Backend Engine**: Seamlessly executes either PyTorch (`.pth`/`.pt`) or ONNX (`.onnx`) models using the same standard interface.
+- **Sliding-Window Inference**: Efficiently process ultra-high-resolution images by dividing them into overlapping patches.
+- **Gaussian & Average Blending**: Reconstructs the full image from patches using overlapping Gaussian blending filters to eliminate edge-seam artifacts.
+- **Test-Time Augmentation (TTA)**: Performs multi-way augmentations (original, horizontal flip, vertical flip, and rotations) to produce highly robust prediction masks.
+- **Validation Metrics**: Compute standard segmentation metrics like IoU, Dice Coefficient, Precision, Recall, and Pixel Accuracy.
+
+---
+
+## Installation
+
+You can install `findcrack` directly from source or via PyPI (once published):
+
+```bash
+# Install via pip
+pip install findcrack
+
+# Or using uv
+uv add findcrack
+```
+
+---
+
+## Quickstart
+
+Here is how to load a pre-trained model and run crack detection on a large image:
+
+```python
+import cv2
+from findcrack import CrackInferencePipeline, load_model
+
+# 1. Load a pre-trained model from the official registry (or use your own URL)
+# The weights are downloaded dynamically from GitHub releases on first use.
+model = load_model("Seg_UNET_CFD_actual_v1", device="cuda")
+
+# 2. Setup the inference pipeline
+pipeline = CrackInferencePipeline(
+    model=model,
+    device="cuda",
+    patch_size=512,
+    overlap_ratio=0.2,
+    confidence_threhold=0.5,
+    use_tta=True  # Enables multi-way Test-Time Augmentation
+)
+
+# 3. Perform inference
+results = pipeline.predict("path/to/high_res_concrete.jpg")
+
+# The results dictionary contains:
+# - results["original_image"]: Original RGB image (numpy array)
+# - results["confidence_map"]: Float probability map [0.0 - 1.0]
+# - results["binary_mask"]: Binary segmentation mask [0 or 255]
+
+# Save the output mask
+cv2.imwrite("detected_cracks.png", results["binary_mask"])
+```
+
+---
+
+## API Reference
+
+### Model Loading & Caching
+
+#### `load_model(variant: str, device: str = "cpu", force_download: bool = False, architecture = None, **kwargs)`
+Loads a model variant from the local registry or directly from a remote HTTP(S) URL.
+
+- **Parameters**:
+  - `variant`: The name of a registered variant (e.g., `"Seg_UNET_CFD_actual_v1"`) or a direct HTTP(S) URL to a weights file.
+  - `device`: Target execution device (`"cpu"`, `"cuda"`, or `"mps"`).
+  - `force_download`: If `True`, re-downloads weights even if cached locally.
+  - `architecture`: PyTorch architecture class (e.g., `UNet`, `DeepCrack`) - required only if loading a raw `.pth`/`.pt` file from a custom URL.
+
+```python
+from findcrack import load_model, UNet
+
+# Load custom model weights directly from an external URL
+model = load_model(
+    variant="https://my-domain.com/custom_unet.pth",
+    architecture=UNet,
+    device="cuda"
+)
+```
+
+#### `list_models()`
+Returns a list of all pre-trained models available in the built-in registry.
+
+#### `register_model(name: str, url: str, architecture = None, kwargs: dict = None, backend: str = "pytorch")`
+Registers a custom variant dynamically at runtime.
+
+---
+
+### Pipeline Configuration
+
+#### `CrackInferencePipeline(model, device: str = "cuda", patch_size: int = 512, overlap_ratio: float = 0.2, confidence_threhold: float = 0.5, use_tta: bool = False)`
+Handles sliding window preprocessing, execution, TTA, and patching reconstruction.
+
+---
+
+## Directory Structure
+
+```text
+src/
+└── findcrack/
+    ├── __init__.py          # Main API endpoints (load_model, CrackInferencePipeline, etc.)
+    ├── metrics.py           # Segmentation evaluation metrics (IoU, Dice, etc.)
+    ├── patching.py          # Sliding window extraction and blend reconstruction
+    ├── pipeline.py          # Crack Inference Pipeline wrapper
+    ├── preprocess.py        # Color-space CLAHE contrast enhancement & transforms
+    ├── tta.py               # Test-Time Augmentation forward pass routines
+    └── models/
+        ├── __init__.py      # Model module exports
+        ├── unet.py          # U-Net model definition
+        ├── deepcrack.py     # DeepCrack model definition
+        ├── onnx_wrapper.py  # Wrapper for running ONNX models as nn.Modules
+        └── zoo.py           # Remote weight registry and cached loaders
+```
+
+---
+
+## License
+
+This project is licensed under the MIT License.

findcrack-0.2.0.dist-info/RECORD

@@ -0,0 +1,15 @@
+findcrack/__init__.py,sha256=IXBxyQAFE3kg0tPLIfbEHMzf1p928jA4zfllDN-UKlQ,405
+findcrack/metrics.py,sha256=GFN34S165ol2Aty3QMAzHLLArfa9AxD3zT3bnvi-rmY,805
+findcrack/patching.py,sha256=DAvtGborNiLYMb0-CEIDfH_dwAwZG1QIURmUUkHfyW0,6807
+findcrack/pipeline.py,sha256=uZTNybTXGp04PlAO28dcvll9jyFrgd25yDNVUcsaK6g,3319
+findcrack/preprocess.py,sha256=xUqKD90DdbiuIlAj5Gh3b0uEGhD0mAqTV6nUXuBha0s,879
+findcrack/tta.py,sha256=tihzsJ1E-CQmjToHrRhMsKgEIH97JeU-5K-HM3b9JsE,1437
+findcrack/models/__init__.py,sha256=t6ZiIRBXzE-cI8_fyJDzcQp0iZorF-Zreib3R4uyJkg,323
+findcrack/models/deepcrack.py,sha256=8pT5dqFYDBw3afcy15SeTlF6DKu_H7o1-0qddZS_Tdg,3817
+findcrack/models/onnx_wrapper.py,sha256=GZQlXRMQtouOVV9u-L0DJmbFLgPnYB2NQ6NaV9EoOPQ,1531
+findcrack/models/unet.py,sha256=0ROuw4x44OirBdP1j8q3l7s3FftuMYaDKs3CBFNudrc,3573
+findcrack/models/zoo.py,sha256=rPiSJDJc--AIMR5_IaD1eGTNe-P4qsEl-3Vjm4Fiqz4,7032
+findcrack-0.2.0.dist-info/METADATA,sha256=EzmmjFKqbCPkrzCuIc5YfROL3BPNbqR595gUzeh9AJk,5803
+findcrack-0.2.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+findcrack-0.2.0.dist-info/top_level.txt,sha256=shEfebjmHLz_MFykmip3VjRyqgErkbOp3Ov5WfYDaWw,10
+findcrack-0.2.0.dist-info/RECORD,,

findcrack-0.2.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

findcrack-0.2.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+findcrack