findcrack 0.2.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- findcrack-0.2.0/PKG-INFO +152 -0
- findcrack-0.2.0/README.md +130 -0
- findcrack-0.2.0/pyproject.toml +36 -0
- findcrack-0.2.0/setup.cfg +4 -0
- findcrack-0.2.0/src/findcrack/__init__.py +17 -0
- findcrack-0.2.0/src/findcrack/metrics.py +23 -0
- findcrack-0.2.0/src/findcrack/models/__init__.py +16 -0
- findcrack-0.2.0/src/findcrack/models/deepcrack.py +93 -0
- findcrack-0.2.0/src/findcrack/models/onnx_wrapper.py +35 -0
- findcrack-0.2.0/src/findcrack/models/unet.py +103 -0
- findcrack-0.2.0/src/findcrack/models/zoo.py +192 -0
- findcrack-0.2.0/src/findcrack/patching.py +172 -0
- findcrack-0.2.0/src/findcrack/pipeline.py +89 -0
- findcrack-0.2.0/src/findcrack/preprocess.py +25 -0
- findcrack-0.2.0/src/findcrack/tta.py +40 -0
- findcrack-0.2.0/src/findcrack.egg-info/PKG-INFO +152 -0
- findcrack-0.2.0/src/findcrack.egg-info/SOURCES.txt +18 -0
- findcrack-0.2.0/src/findcrack.egg-info/dependency_links.txt +1 -0
- findcrack-0.2.0/src/findcrack.egg-info/requires.txt +6 -0
- findcrack-0.2.0/src/findcrack.egg-info/top_level.txt +1 -0
findcrack-0.2.0/PKG-INFO
ADDED
|
@@ -0,0 +1,152 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: findcrack
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: A deep learning crack detection package supporting U-Net and DeepCrack models with PyTorch and ONNX backends.
|
|
5
|
+
Author: StrikerEurika
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/StrikerEurika/findcrack
|
|
8
|
+
Project-URL: Repository, https://github.com/StrikerEurika/findcrack
|
|
9
|
+
Project-URL: Releases, https://github.com/StrikerEurika/findcrack/releases
|
|
10
|
+
Classifier: Programming Language :: Python :: 3
|
|
11
|
+
Classifier: Operating System :: OS Independent
|
|
12
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
13
|
+
Classifier: Topic :: Scientific/Engineering :: Image Processing
|
|
14
|
+
Requires-Python: >=3.11
|
|
15
|
+
Description-Content-Type: text/markdown
|
|
16
|
+
Requires-Dist: albumentations>=2.0.8
|
|
17
|
+
Requires-Dist: onnxruntime>=1.27.0
|
|
18
|
+
Requires-Dist: opencv-python>=4.13.0.92
|
|
19
|
+
Requires-Dist: torch>=2.6.0
|
|
20
|
+
Requires-Dist: torchaudio>=2.6.0
|
|
21
|
+
Requires-Dist: torchvision>=0.21.0
|
|
22
|
+
|
|
23
|
+
# findcrack
|
|
24
|
+
|
|
25
|
+
`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).
|
|
26
|
+
|
|
27
|
+
---
|
|
28
|
+
|
|
29
|
+
## Features
|
|
30
|
+
|
|
31
|
+
- **Pre-trained Model Zoo**: Fetch pre-trained model weights (e.g., `Seg_UNET_CFD_actual_v1`, `Seg_UNET_CFD_actual_v2`) dynamically on demand.
|
|
32
|
+
- **Unified Backend Engine**: Seamlessly executes either PyTorch (`.pth`/`.pt`) or ONNX (`.onnx`) models using the same standard interface.
|
|
33
|
+
- **Sliding-Window Inference**: Efficiently process ultra-high-resolution images by dividing them into overlapping patches.
|
|
34
|
+
- **Gaussian & Average Blending**: Reconstructs the full image from patches using overlapping Gaussian blending filters to eliminate edge-seam artifacts.
|
|
35
|
+
- **Test-Time Augmentation (TTA)**: Performs multi-way augmentations (original, horizontal flip, vertical flip, and rotations) to produce highly robust prediction masks.
|
|
36
|
+
- **Validation Metrics**: Compute standard segmentation metrics like IoU, Dice Coefficient, Precision, Recall, and Pixel Accuracy.
|
|
37
|
+
|
|
38
|
+
---
|
|
39
|
+
|
|
40
|
+
## Installation
|
|
41
|
+
|
|
42
|
+
You can install `findcrack` directly from source or via PyPI (once published):
|
|
43
|
+
|
|
44
|
+
```bash
|
|
45
|
+
# Install via pip
|
|
46
|
+
pip install findcrack
|
|
47
|
+
|
|
48
|
+
# Or using uv
|
|
49
|
+
uv add findcrack
|
|
50
|
+
```
|
|
51
|
+
|
|
52
|
+
---
|
|
53
|
+
|
|
54
|
+
## Quickstart
|
|
55
|
+
|
|
56
|
+
Here is how to load a pre-trained model and run crack detection on a large image:
|
|
57
|
+
|
|
58
|
+
```python
|
|
59
|
+
import cv2
|
|
60
|
+
from findcrack import CrackInferencePipeline, load_model
|
|
61
|
+
|
|
62
|
+
# 1. Load a pre-trained model from the official registry (or use your own URL)
|
|
63
|
+
# The weights are downloaded dynamically from GitHub releases on first use.
|
|
64
|
+
model = load_model("Seg_UNET_CFD_actual_v1", device="cuda")
|
|
65
|
+
|
|
66
|
+
# 2. Setup the inference pipeline
|
|
67
|
+
pipeline = CrackInferencePipeline(
|
|
68
|
+
model=model,
|
|
69
|
+
device="cuda",
|
|
70
|
+
patch_size=512,
|
|
71
|
+
overlap_ratio=0.2,
|
|
72
|
+
confidence_threhold=0.5,
|
|
73
|
+
use_tta=True # Enables multi-way Test-Time Augmentation
|
|
74
|
+
)
|
|
75
|
+
|
|
76
|
+
# 3. Perform inference
|
|
77
|
+
results = pipeline.predict("path/to/high_res_concrete.jpg")
|
|
78
|
+
|
|
79
|
+
# The results dictionary contains:
|
|
80
|
+
# - results["original_image"]: Original RGB image (numpy array)
|
|
81
|
+
# - results["confidence_map"]: Float probability map [0.0 - 1.0]
|
|
82
|
+
# - results["binary_mask"]: Binary segmentation mask [0 or 255]
|
|
83
|
+
|
|
84
|
+
# Save the output mask
|
|
85
|
+
cv2.imwrite("detected_cracks.png", results["binary_mask"])
|
|
86
|
+
```
|
|
87
|
+
|
|
88
|
+
---
|
|
89
|
+
|
|
90
|
+
## API Reference
|
|
91
|
+
|
|
92
|
+
### Model Loading & Caching
|
|
93
|
+
|
|
94
|
+
#### `load_model(variant: str, device: str = "cpu", force_download: bool = False, architecture = None, **kwargs)`
|
|
95
|
+
Loads a model variant from the local registry or directly from a remote HTTP(S) URL.
|
|
96
|
+
|
|
97
|
+
- **Parameters**:
|
|
98
|
+
- `variant`: The name of a registered variant (e.g., `"Seg_UNET_CFD_actual_v1"`) or a direct HTTP(S) URL to a weights file.
|
|
99
|
+
- `device`: Target execution device (`"cpu"`, `"cuda"`, or `"mps"`).
|
|
100
|
+
- `force_download`: If `True`, re-downloads weights even if cached locally.
|
|
101
|
+
- `architecture`: PyTorch architecture class (e.g., `UNet`, `DeepCrack`) - required only if loading a raw `.pth`/`.pt` file from a custom URL.
|
|
102
|
+
|
|
103
|
+
```python
|
|
104
|
+
from findcrack import load_model, UNet
|
|
105
|
+
|
|
106
|
+
# Load custom model weights directly from an external URL
|
|
107
|
+
model = load_model(
|
|
108
|
+
variant="https://my-domain.com/custom_unet.pth",
|
|
109
|
+
architecture=UNet,
|
|
110
|
+
device="cuda"
|
|
111
|
+
)
|
|
112
|
+
```
|
|
113
|
+
|
|
114
|
+
#### `list_models()`
|
|
115
|
+
Returns a list of all pre-trained models available in the built-in registry.
|
|
116
|
+
|
|
117
|
+
#### `register_model(name: str, url: str, architecture = None, kwargs: dict = None, backend: str = "pytorch")`
|
|
118
|
+
Registers a custom variant dynamically at runtime.
|
|
119
|
+
|
|
120
|
+
---
|
|
121
|
+
|
|
122
|
+
### Pipeline Configuration
|
|
123
|
+
|
|
124
|
+
#### `CrackInferencePipeline(model, device: str = "cuda", patch_size: int = 512, overlap_ratio: float = 0.2, confidence_threhold: float = 0.5, use_tta: bool = False)`
|
|
125
|
+
Handles sliding window preprocessing, execution, TTA, and patching reconstruction.
|
|
126
|
+
|
|
127
|
+
---
|
|
128
|
+
|
|
129
|
+
## Directory Structure
|
|
130
|
+
|
|
131
|
+
```text
|
|
132
|
+
src/
|
|
133
|
+
└── findcrack/
|
|
134
|
+
├── __init__.py # Main API endpoints (load_model, CrackInferencePipeline, etc.)
|
|
135
|
+
├── metrics.py # Segmentation evaluation metrics (IoU, Dice, etc.)
|
|
136
|
+
├── patching.py # Sliding window extraction and blend reconstruction
|
|
137
|
+
├── pipeline.py # Crack Inference Pipeline wrapper
|
|
138
|
+
├── preprocess.py # Color-space CLAHE contrast enhancement & transforms
|
|
139
|
+
├── tta.py # Test-Time Augmentation forward pass routines
|
|
140
|
+
└── models/
|
|
141
|
+
├── __init__.py # Model module exports
|
|
142
|
+
├── unet.py # U-Net model definition
|
|
143
|
+
├── deepcrack.py # DeepCrack model definition
|
|
144
|
+
├── onnx_wrapper.py # Wrapper for running ONNX models as nn.Modules
|
|
145
|
+
└── zoo.py # Remote weight registry and cached loaders
|
|
146
|
+
```
|
|
147
|
+
|
|
148
|
+
---
|
|
149
|
+
|
|
150
|
+
## License
|
|
151
|
+
|
|
152
|
+
This project is licensed under the MIT License.
|
|
@@ -0,0 +1,130 @@
|
|
|
1
|
+
# findcrack
|
|
2
|
+
|
|
3
|
+
`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).
|
|
4
|
+
|
|
5
|
+
---
|
|
6
|
+
|
|
7
|
+
## Features
|
|
8
|
+
|
|
9
|
+
- **Pre-trained Model Zoo**: Fetch pre-trained model weights (e.g., `Seg_UNET_CFD_actual_v1`, `Seg_UNET_CFD_actual_v2`) dynamically on demand.
|
|
10
|
+
- **Unified Backend Engine**: Seamlessly executes either PyTorch (`.pth`/`.pt`) or ONNX (`.onnx`) models using the same standard interface.
|
|
11
|
+
- **Sliding-Window Inference**: Efficiently process ultra-high-resolution images by dividing them into overlapping patches.
|
|
12
|
+
- **Gaussian & Average Blending**: Reconstructs the full image from patches using overlapping Gaussian blending filters to eliminate edge-seam artifacts.
|
|
13
|
+
- **Test-Time Augmentation (TTA)**: Performs multi-way augmentations (original, horizontal flip, vertical flip, and rotations) to produce highly robust prediction masks.
|
|
14
|
+
- **Validation Metrics**: Compute standard segmentation metrics like IoU, Dice Coefficient, Precision, Recall, and Pixel Accuracy.
|
|
15
|
+
|
|
16
|
+
---
|
|
17
|
+
|
|
18
|
+
## Installation
|
|
19
|
+
|
|
20
|
+
You can install `findcrack` directly from source or via PyPI (once published):
|
|
21
|
+
|
|
22
|
+
```bash
|
|
23
|
+
# Install via pip
|
|
24
|
+
pip install findcrack
|
|
25
|
+
|
|
26
|
+
# Or using uv
|
|
27
|
+
uv add findcrack
|
|
28
|
+
```
|
|
29
|
+
|
|
30
|
+
---
|
|
31
|
+
|
|
32
|
+
## Quickstart
|
|
33
|
+
|
|
34
|
+
Here is how to load a pre-trained model and run crack detection on a large image:
|
|
35
|
+
|
|
36
|
+
```python
|
|
37
|
+
import cv2
|
|
38
|
+
from findcrack import CrackInferencePipeline, load_model
|
|
39
|
+
|
|
40
|
+
# 1. Load a pre-trained model from the official registry (or use your own URL)
|
|
41
|
+
# The weights are downloaded dynamically from GitHub releases on first use.
|
|
42
|
+
model = load_model("Seg_UNET_CFD_actual_v1", device="cuda")
|
|
43
|
+
|
|
44
|
+
# 2. Setup the inference pipeline
|
|
45
|
+
pipeline = CrackInferencePipeline(
|
|
46
|
+
model=model,
|
|
47
|
+
device="cuda",
|
|
48
|
+
patch_size=512,
|
|
49
|
+
overlap_ratio=0.2,
|
|
50
|
+
confidence_threhold=0.5,
|
|
51
|
+
use_tta=True # Enables multi-way Test-Time Augmentation
|
|
52
|
+
)
|
|
53
|
+
|
|
54
|
+
# 3. Perform inference
|
|
55
|
+
results = pipeline.predict("path/to/high_res_concrete.jpg")
|
|
56
|
+
|
|
57
|
+
# The results dictionary contains:
|
|
58
|
+
# - results["original_image"]: Original RGB image (numpy array)
|
|
59
|
+
# - results["confidence_map"]: Float probability map [0.0 - 1.0]
|
|
60
|
+
# - results["binary_mask"]: Binary segmentation mask [0 or 255]
|
|
61
|
+
|
|
62
|
+
# Save the output mask
|
|
63
|
+
cv2.imwrite("detected_cracks.png", results["binary_mask"])
|
|
64
|
+
```
|
|
65
|
+
|
|
66
|
+
---
|
|
67
|
+
|
|
68
|
+
## API Reference
|
|
69
|
+
|
|
70
|
+
### Model Loading & Caching
|
|
71
|
+
|
|
72
|
+
#### `load_model(variant: str, device: str = "cpu", force_download: bool = False, architecture = None, **kwargs)`
|
|
73
|
+
Loads a model variant from the local registry or directly from a remote HTTP(S) URL.
|
|
74
|
+
|
|
75
|
+
- **Parameters**:
|
|
76
|
+
- `variant`: The name of a registered variant (e.g., `"Seg_UNET_CFD_actual_v1"`) or a direct HTTP(S) URL to a weights file.
|
|
77
|
+
- `device`: Target execution device (`"cpu"`, `"cuda"`, or `"mps"`).
|
|
78
|
+
- `force_download`: If `True`, re-downloads weights even if cached locally.
|
|
79
|
+
- `architecture`: PyTorch architecture class (e.g., `UNet`, `DeepCrack`) - required only if loading a raw `.pth`/`.pt` file from a custom URL.
|
|
80
|
+
|
|
81
|
+
```python
|
|
82
|
+
from findcrack import load_model, UNet
|
|
83
|
+
|
|
84
|
+
# Load custom model weights directly from an external URL
|
|
85
|
+
model = load_model(
|
|
86
|
+
variant="https://my-domain.com/custom_unet.pth",
|
|
87
|
+
architecture=UNet,
|
|
88
|
+
device="cuda"
|
|
89
|
+
)
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
#### `list_models()`
|
|
93
|
+
Returns a list of all pre-trained models available in the built-in registry.
|
|
94
|
+
|
|
95
|
+
#### `register_model(name: str, url: str, architecture = None, kwargs: dict = None, backend: str = "pytorch")`
|
|
96
|
+
Registers a custom variant dynamically at runtime.
|
|
97
|
+
|
|
98
|
+
---
|
|
99
|
+
|
|
100
|
+
### Pipeline Configuration
|
|
101
|
+
|
|
102
|
+
#### `CrackInferencePipeline(model, device: str = "cuda", patch_size: int = 512, overlap_ratio: float = 0.2, confidence_threhold: float = 0.5, use_tta: bool = False)`
|
|
103
|
+
Handles sliding window preprocessing, execution, TTA, and patching reconstruction.
|
|
104
|
+
|
|
105
|
+
---
|
|
106
|
+
|
|
107
|
+
## Directory Structure
|
|
108
|
+
|
|
109
|
+
```text
|
|
110
|
+
src/
|
|
111
|
+
└── findcrack/
|
|
112
|
+
├── __init__.py # Main API endpoints (load_model, CrackInferencePipeline, etc.)
|
|
113
|
+
├── metrics.py # Segmentation evaluation metrics (IoU, Dice, etc.)
|
|
114
|
+
├── patching.py # Sliding window extraction and blend reconstruction
|
|
115
|
+
├── pipeline.py # Crack Inference Pipeline wrapper
|
|
116
|
+
├── preprocess.py # Color-space CLAHE contrast enhancement & transforms
|
|
117
|
+
├── tta.py # Test-Time Augmentation forward pass routines
|
|
118
|
+
└── models/
|
|
119
|
+
├── __init__.py # Model module exports
|
|
120
|
+
├── unet.py # U-Net model definition
|
|
121
|
+
├── deepcrack.py # DeepCrack model definition
|
|
122
|
+
├── onnx_wrapper.py # Wrapper for running ONNX models as nn.Modules
|
|
123
|
+
└── zoo.py # Remote weight registry and cached loaders
|
|
124
|
+
```
|
|
125
|
+
|
|
126
|
+
---
|
|
127
|
+
|
|
128
|
+
## License
|
|
129
|
+
|
|
130
|
+
This project is licensed under the MIT License.
|
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "findcrack"
|
|
3
|
+
version = "0.2.0"
|
|
4
|
+
description = "A deep learning crack detection package supporting U-Net and DeepCrack models with PyTorch and ONNX backends."
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.11"
|
|
7
|
+
authors = [
|
|
8
|
+
{ name = "StrikerEurika" }
|
|
9
|
+
]
|
|
10
|
+
license = "MIT"
|
|
11
|
+
classifiers = [
|
|
12
|
+
"Programming Language :: Python :: 3",
|
|
13
|
+
"Operating System :: OS Independent",
|
|
14
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
15
|
+
"Topic :: Scientific/Engineering :: Image Processing",
|
|
16
|
+
]
|
|
17
|
+
dependencies = [
|
|
18
|
+
"albumentations>=2.0.8",
|
|
19
|
+
"onnxruntime>=1.27.0",
|
|
20
|
+
"opencv-python>=4.13.0.92",
|
|
21
|
+
"torch>=2.6.0",
|
|
22
|
+
"torchaudio>=2.6.0",
|
|
23
|
+
"torchvision>=0.21.0",
|
|
24
|
+
]
|
|
25
|
+
|
|
26
|
+
[project.urls]
|
|
27
|
+
Homepage = "https://github.com/StrikerEurika/findcrack"
|
|
28
|
+
Repository = "https://github.com/StrikerEurika/findcrack"
|
|
29
|
+
Releases = "https://github.com/StrikerEurika/findcrack/releases"
|
|
30
|
+
|
|
31
|
+
[tool.setuptools.packages.find]
|
|
32
|
+
where = ["src"]
|
|
33
|
+
|
|
34
|
+
[build-system]
|
|
35
|
+
requires = ["setuptools>=61.0", "wheel"]
|
|
36
|
+
build-backend = "setuptools.build_meta"
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
from .pipeline import CrackInferencePipeline
|
|
2
|
+
from .models import load_model, UNet, DeepCrack, list_models, register_model
|
|
3
|
+
from .metrics import calculate_metrics
|
|
4
|
+
from .preprocess import apply_lab_clahe
|
|
5
|
+
|
|
6
|
+
__version__ = "0.1.0"
|
|
7
|
+
|
|
8
|
+
__all__ = [
|
|
9
|
+
"CrackInferencePipeline",
|
|
10
|
+
"load_model",
|
|
11
|
+
"UNet",
|
|
12
|
+
"DeepCrack",
|
|
13
|
+
"calculate_metrics",
|
|
14
|
+
"apply_lab_clahe",
|
|
15
|
+
"list_models",
|
|
16
|
+
"register_model",
|
|
17
|
+
]
|
|
@@ -0,0 +1,23 @@
|
|
|
1
|
+
import numpy as np
|
|
2
|
+
|
|
3
|
+
def calculate_metrics(y_true: np.ndarray, y_prediction: np.ndarray, epsilon: float = 1e-7) -> dict:
|
|
4
|
+
"""
|
|
5
|
+
Calculates IoU, Dice, Precision, Recall, and Pixel Accuracy.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
y_true = y_true.astype(bool)
|
|
9
|
+
y_prediction = y_prediction.astype(bool)
|
|
10
|
+
|
|
11
|
+
# True Positives, False Positives, True Negatives, False Negatives
|
|
12
|
+
TP = np.sum(y_true & y_prediction)
|
|
13
|
+
FP = np.sum(~y_true & y_prediction)
|
|
14
|
+
FN = np.sum(y_true & ~y_prediction)
|
|
15
|
+
TN = np.sum(~y_true & ~y_prediction)
|
|
16
|
+
|
|
17
|
+
return {
|
|
18
|
+
"IoU": TP / (TP + FP + FN + epsilon),
|
|
19
|
+
"Dice": (2 * TP) / (2 * TP + FP + FN + epsilon),
|
|
20
|
+
"Precision": TP / (TP + FP + epsilon),
|
|
21
|
+
"Recall": TP / (TP + FN + epsilon),
|
|
22
|
+
"Pixel Accuracy": (TP + TN) / (TP + TN + FP + FN + epsilon)
|
|
23
|
+
}
|
|
@@ -0,0 +1,16 @@
|
|
|
1
|
+
from .unet import UNet
|
|
2
|
+
from .deepcrack import DeepCrack
|
|
3
|
+
from .onnx_wrapper import ONNXModelWrapper
|
|
4
|
+
from .zoo import load_model, MODEL_REGISTRY, list_models, register_model
|
|
5
|
+
|
|
6
|
+
__all__ = [
|
|
7
|
+
"UNet",
|
|
8
|
+
"DeepCrack",
|
|
9
|
+
"ONNXModelWrapper",
|
|
10
|
+
"load_model",
|
|
11
|
+
"MODEL_REGISTRY",
|
|
12
|
+
"list_models",
|
|
13
|
+
"register_model",
|
|
14
|
+
]
|
|
15
|
+
|
|
16
|
+
|
|
@@ -0,0 +1,93 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
import torch.nn as nn
|
|
3
|
+
import torch.nn.functional as F
|
|
4
|
+
|
|
5
|
+
class DoubleConv(nn.Module):
|
|
6
|
+
"""(convolution => [BN] => ReLU) * 2"""
|
|
7
|
+
def __init__(self, in_channels, out_channels):
|
|
8
|
+
super().__init__()
|
|
9
|
+
self.double_conv = nn.Sequential(
|
|
10
|
+
nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
|
|
11
|
+
nn.BatchNorm2d(out_channels),
|
|
12
|
+
nn.ReLU(inplace=True),
|
|
13
|
+
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),
|
|
14
|
+
nn.BatchNorm2d(out_channels),
|
|
15
|
+
nn.ReLU(inplace=True)
|
|
16
|
+
)
|
|
17
|
+
|
|
18
|
+
def forward(self, x):
|
|
19
|
+
return self.double_conv(x)
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
class DeepCrack(nn.Module):
|
|
23
|
+
"""
|
|
24
|
+
DeepCrack: A Deep Hierarchical Feature Learning Architecture for Crack Segmentation.
|
|
25
|
+
Fuses hierarchical convolutional features from both the encoder and decoder stages
|
|
26
|
+
at the same scale.
|
|
27
|
+
"""
|
|
28
|
+
def __init__(self, n_channels: int = 3, n_classes: int = 1):
|
|
29
|
+
super().__init__()
|
|
30
|
+
self.n_channels = n_channels
|
|
31
|
+
self.n_classes = n_classes
|
|
32
|
+
|
|
33
|
+
# Encoder (downsampling blocks)
|
|
34
|
+
self.enc1 = DoubleConv(n_channels, 64)
|
|
35
|
+
self.enc2 = DoubleConv(64, 128)
|
|
36
|
+
self.enc3 = DoubleConv(128, 256)
|
|
37
|
+
self.enc4 = DoubleConv(256, 512)
|
|
38
|
+
self.enc5 = DoubleConv(512, 512)
|
|
39
|
+
|
|
40
|
+
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
|
|
41
|
+
|
|
42
|
+
# Decoder (upsampling & concatenation blocks)
|
|
43
|
+
self.dec5 = DoubleConv(512, 512)
|
|
44
|
+
self.dec4 = DoubleConv(512 + 512, 256)
|
|
45
|
+
self.dec3 = DoubleConv(256 + 256, 128)
|
|
46
|
+
self.dec2 = DoubleConv(128 + 128, 64)
|
|
47
|
+
self.dec1 = DoubleConv(64 + 64, 64)
|
|
48
|
+
|
|
49
|
+
# Side prediction layers (maps feature channels to class channels at each scale)
|
|
50
|
+
self.side1 = nn.Conv2d(64, n_classes, kernel_size=1)
|
|
51
|
+
self.side2 = nn.Conv2d(64, n_classes, kernel_size=1)
|
|
52
|
+
self.side3 = nn.Conv2d(128, n_classes, kernel_size=1)
|
|
53
|
+
self.side4 = nn.Conv2d(256, n_classes, kernel_size=1)
|
|
54
|
+
self.side5 = nn.Conv2d(512, n_classes, kernel_size=1)
|
|
55
|
+
|
|
56
|
+
# Fusion layer that combines all 5 side predictions into the final output
|
|
57
|
+
self.fuse = nn.Conv2d(n_classes * 5, n_classes, kernel_size=1)
|
|
58
|
+
|
|
59
|
+
def forward(self, x):
|
|
60
|
+
# 1. Encoder path
|
|
61
|
+
e1 = self.enc1(x)
|
|
62
|
+
e2 = self.enc2(self.pool(e1))
|
|
63
|
+
e3 = self.enc3(self.pool(e2))
|
|
64
|
+
e4 = self.enc4(self.pool(e3))
|
|
65
|
+
e5 = self.enc5(self.pool(e4))
|
|
66
|
+
|
|
67
|
+
# 2. Decoder path (with bilinear interpolation upsampling)
|
|
68
|
+
d5 = self.dec5(e5)
|
|
69
|
+
|
|
70
|
+
d4_up = F.interpolate(d5, size=e4.shape[2:], mode='bilinear', align_corners=True)
|
|
71
|
+
d4 = self.dec4(torch.cat([d4_up, e4], dim=1))
|
|
72
|
+
|
|
73
|
+
d3_up = F.interpolate(d4, size=e3.shape[2:], mode='bilinear', align_corners=True)
|
|
74
|
+
d3 = self.dec3(torch.cat([d3_up, e3], dim=1))
|
|
75
|
+
|
|
76
|
+
d2_up = F.interpolate(d3, size=e2.shape[2:], mode='bilinear', align_corners=True)
|
|
77
|
+
d2 = self.dec2(torch.cat([d2_up, e2], dim=1))
|
|
78
|
+
|
|
79
|
+
d1_up = F.interpolate(d2, size=e1.shape[2:], mode='bilinear', align_corners=True)
|
|
80
|
+
d1 = self.dec1(torch.cat([d1_up, e1], dim=1))
|
|
81
|
+
|
|
82
|
+
# 3. Extract side predictions and upsample to input dimensions
|
|
83
|
+
h, w = x.shape[2:]
|
|
84
|
+
s1 = F.interpolate(self.side1(d1), size=(h, w), mode='bilinear', align_corners=True)
|
|
85
|
+
s2 = F.interpolate(self.side2(d2), size=(h, w), mode='bilinear', align_corners=True)
|
|
86
|
+
s3 = F.interpolate(self.side3(d3), size=(h, w), mode='bilinear', align_corners=True)
|
|
87
|
+
s4 = F.interpolate(self.side4(d4), size=(h, w), mode='bilinear', align_corners=True)
|
|
88
|
+
s5 = F.interpolate(self.side5(d5), size=(h, w), mode='bilinear', align_corners=True)
|
|
89
|
+
|
|
90
|
+
# 4. Fuse side predictions
|
|
91
|
+
fused = self.fuse(torch.cat([s1, s2, s3, s4, s5], dim=1))
|
|
92
|
+
|
|
93
|
+
return fused
|
|
@@ -0,0 +1,35 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
import torch.nn as nn
|
|
3
|
+
import numpy as np
|
|
4
|
+
|
|
5
|
+
class ONNXModelWrapper(nn.Module):
|
|
6
|
+
"""
|
|
7
|
+
Wraps an ONNX Runtime InferenceSession inside a PyTorch nn.Module.
|
|
8
|
+
This allows running inference on ONNX models using the exact same code
|
|
9
|
+
and APIs as PyTorch models, supporting both CPU/GPU tensor operations
|
|
10
|
+
and test-time augmentation (TTA) pipelines.
|
|
11
|
+
"""
|
|
12
|
+
def __init__(self, model_path: str, device: str = "cpu"):
|
|
13
|
+
super().__init__()
|
|
14
|
+
import onnxruntime as ort
|
|
15
|
+
|
|
16
|
+
# Select execution providers based on the target device
|
|
17
|
+
if device == "cuda" or (isinstance(device, torch.device) and device.type == "cuda"):
|
|
18
|
+
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
|
|
19
|
+
else:
|
|
20
|
+
providers = ["CPUExecutionProvider"]
|
|
21
|
+
|
|
22
|
+
self.session = ort.InferenceSession(model_path, providers=providers)
|
|
23
|
+
self.input_name = self.session.get_inputs()[0].name
|
|
24
|
+
self.output_name = self.session.get_outputs()[0].name
|
|
25
|
+
|
|
26
|
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
27
|
+
# 1. Convert PyTorch tensor to NumPy array (typically expected as float32)
|
|
28
|
+
x_np = x.detach().cpu().numpy().astype(np.float32)
|
|
29
|
+
|
|
30
|
+
# 2. Run inference using ONNX Runtime
|
|
31
|
+
outputs = self.session.run([self.output_name], {self.input_name: x_np})
|
|
32
|
+
|
|
33
|
+
# 3. Convert prediction back to PyTorch tensor and move to the original device
|
|
34
|
+
out_tensor = torch.from_numpy(outputs[0]).to(x.device)
|
|
35
|
+
return out_tensor
|
|
@@ -0,0 +1,103 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
import torch.nn as nn
|
|
3
|
+
|
|
4
|
+
# U-Net Model Definition
|
|
5
|
+
class DoubleConv(nn.Module):
|
|
6
|
+
""" (convolution => [BN] => ReLU) * 2 """
|
|
7
|
+
def __init__(self, in_channels, out_channels, mid_channels=None):
|
|
8
|
+
super().__init__()
|
|
9
|
+
if not mid_channels:
|
|
10
|
+
mid_channels = out_channels
|
|
11
|
+
self.double_conv = nn.Sequential(
|
|
12
|
+
nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
|
|
13
|
+
nn.BatchNorm2d(mid_channels),
|
|
14
|
+
nn.ReLU(inplace=True),
|
|
15
|
+
nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
|
|
16
|
+
nn.BatchNorm2d(out_channels),
|
|
17
|
+
nn.ReLU(inplace=True)
|
|
18
|
+
)
|
|
19
|
+
|
|
20
|
+
def forward(self, x):
|
|
21
|
+
return self.double_conv(x)
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
class Down(nn.Module):
|
|
25
|
+
""" Downscaling with maxpool then double conv """
|
|
26
|
+
def __init__(self, in_channels, out_channels):
|
|
27
|
+
super().__init__()
|
|
28
|
+
self.maxpool_conv = nn.Sequential(
|
|
29
|
+
nn.MaxPool2d(2),
|
|
30
|
+
DoubleConv(in_channels, out_channels)
|
|
31
|
+
)
|
|
32
|
+
|
|
33
|
+
def forward(self, x):
|
|
34
|
+
return self.maxpool_conv(x)
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
class Up(nn.Module):
|
|
38
|
+
""" Upscaling then double conv """
|
|
39
|
+
def __init__(self, in_channels, out_channels, bilinear=True):
|
|
40
|
+
super().__init__()
|
|
41
|
+
|
|
42
|
+
if bilinear:
|
|
43
|
+
self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
|
|
44
|
+
self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)
|
|
45
|
+
else:
|
|
46
|
+
self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
|
|
47
|
+
self.conv = DoubleConv(in_channels, out_channels)
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
def forward(self, x1, x2):
|
|
51
|
+
x1 = self.up(x1)
|
|
52
|
+
# Input tensor might not be perfectly divisible by 2, so crop x2 to match x1's size
|
|
53
|
+
# This handles potential size mismatches due to padding in earlier layers or odd input dimensions
|
|
54
|
+
diffY = x2.size()[2] - x1.size()[2]
|
|
55
|
+
diffX = x2.size()[3] - x1.size()[3]
|
|
56
|
+
|
|
57
|
+
# Pad x1 if necessary to match x2's size (or crop x2 if it's larger)
|
|
58
|
+
x1 = nn.functional.pad(x1, [diffX // 2, diffX - diffX // 2,
|
|
59
|
+
diffY // 2, diffY - diffY // 2])
|
|
60
|
+
x = torch.cat([x2, x1], dim=1)
|
|
61
|
+
return self.conv(x)
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
class OutConv(nn.Module):
|
|
65
|
+
def __init__(self, in_channels, out_channels):
|
|
66
|
+
super(OutConv, self).__init__()
|
|
67
|
+
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
|
|
68
|
+
|
|
69
|
+
def forward(self, x):
|
|
70
|
+
return self.conv(x)
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
class UNet(nn.Module):
|
|
74
|
+
def __init__(self, n_channels, n_classes, bilinear=False):
|
|
75
|
+
super(UNet, self).__init__()
|
|
76
|
+
self.n_channels = n_channels
|
|
77
|
+
self.n_classes = n_classes
|
|
78
|
+
self.bilinear = bilinear
|
|
79
|
+
|
|
80
|
+
self.inc = DoubleConv(n_channels, 64)
|
|
81
|
+
self.down1 = Down(64, 128)
|
|
82
|
+
self.down2 = Down(128, 256)
|
|
83
|
+
self.down3 = Down(256, 512)
|
|
84
|
+
factor = 2 if bilinear else 1
|
|
85
|
+
self.down4 = Down(512, 1024 // factor)
|
|
86
|
+
self.up1 = Up(1024, 512 // factor, bilinear)
|
|
87
|
+
self.up2 = Up(512, 256 // factor, bilinear)
|
|
88
|
+
self.up3 = Up(256, 128 // factor, bilinear)
|
|
89
|
+
self.up4 = Up(128, 64, bilinear)
|
|
90
|
+
self.outc = OutConv(64, n_classes)
|
|
91
|
+
|
|
92
|
+
def forward(self, x):
|
|
93
|
+
x1 = self.inc(x)
|
|
94
|
+
x2 = self.down1(x1)
|
|
95
|
+
x3 = self.down2(x2)
|
|
96
|
+
x4 = self.down3(x3)
|
|
97
|
+
x5 = self.down4(x4)
|
|
98
|
+
x = self.up1(x5, x4) # F.pad needs to be imported or handled in Up()
|
|
99
|
+
x = self.up2(x, x3)
|
|
100
|
+
x = self.up3(x, x2)
|
|
101
|
+
x = self.up4(x, x1)
|
|
102
|
+
logits = self.outc(x)
|
|
103
|
+
return logits
|
|
@@ -0,0 +1,192 @@
|
|
|
1
|
+
import os
|
|
2
|
+
import torch
|
|
3
|
+
import hashlib
|
|
4
|
+
from torch.hub import get_dir, download_url_to_file
|
|
5
|
+
|
|
6
|
+
from .unet import UNet
|
|
7
|
+
from .deepcrack import DeepCrack
|
|
8
|
+
from .onnx_wrapper import ONNXModelWrapper
|
|
9
|
+
|
|
10
|
+
# Model registry detailing architectures, default arguments, remote URLs, hashes, and backend type.
|
|
11
|
+
MODEL_REGISTRY = {
|
|
12
|
+
"Seg_UNET_CFD_actual_v1": {
|
|
13
|
+
"metadata": {
|
|
14
|
+
"loss_functions": ["TverskyLoss"],
|
|
15
|
+
"input_shape": [3, 512, 512],
|
|
16
|
+
},
|
|
17
|
+
"architecture": UNet,
|
|
18
|
+
"kwargs": {
|
|
19
|
+
"n_channels": 3,
|
|
20
|
+
"n_classes": 1,
|
|
21
|
+
"bilinear": False
|
|
22
|
+
},
|
|
23
|
+
"backend": "pytorch",
|
|
24
|
+
"url": "https://github.com/StrikerEurika/findcrack/releases/download/v0.2.0/Seg_UNET_CFD_actual_v1_best.pth",
|
|
25
|
+
"sha256": None # Users can supply checksums to verify integrity
|
|
26
|
+
},
|
|
27
|
+
"Seg_UNET_CFD_actual_v2": {
|
|
28
|
+
"metadata": {
|
|
29
|
+
"loss_functions": ["BCEWithLogitsLoss", "DiceLoss"],
|
|
30
|
+
"input_shape": [3, 512, 512],
|
|
31
|
+
},
|
|
32
|
+
"architecture": UNet,
|
|
33
|
+
"kwargs": {
|
|
34
|
+
"n_channels": 3,
|
|
35
|
+
"n_classes": 1,
|
|
36
|
+
"bilinear": False
|
|
37
|
+
},
|
|
38
|
+
"backend": "pytorch",
|
|
39
|
+
"url": "https://github.com/StrikerEurika/findcrack/releases/download/v0.2.0/Seg_UNET_CFD_actual_v2_best.pth",
|
|
40
|
+
"sha256": None # Users can supply checksums to verify integrity
|
|
41
|
+
},
|
|
42
|
+
}
|
|
43
|
+
|
|
44
|
+
def verify_sha256(filepath: str, expected_hash: str) -> bool:
|
|
45
|
+
"""Verifies file integrity via SHA256 checksum."""
|
|
46
|
+
if not expected_hash:
|
|
47
|
+
return True
|
|
48
|
+
sha256_hash = hashlib.sha256()
|
|
49
|
+
with open(filepath, "rb") as f:
|
|
50
|
+
for byte_block in iter(lambda: f.read(4096), b""):
|
|
51
|
+
sha256_hash.update(byte_block)
|
|
52
|
+
return sha256_hash.hexdigest() == expected_hash
|
|
53
|
+
|
|
54
|
+
def get_cache_dir() -> str:
|
|
55
|
+
"""Returns the package's weight caching directory."""
|
|
56
|
+
hub_directory = get_dir()
|
|
57
|
+
model_directory = os.path.join(hub_directory, 'checkpoints', 'findcrack')
|
|
58
|
+
os.makedirs(model_directory, exist_ok=True)
|
|
59
|
+
return model_directory
|
|
60
|
+
|
|
61
|
+
def download_weight_file(url: str, cached_file: str, expected_hash: str = None):
|
|
62
|
+
"""Downloads remote weight file and verifies its checksum."""
|
|
63
|
+
print(f"Downloading weights from {url}... (This is a one-time download)")
|
|
64
|
+
try:
|
|
65
|
+
download_url_to_file(url, cached_file, progress=True)
|
|
66
|
+
except Exception as e:
|
|
67
|
+
raise IOError(f"Failed to download weight file from {url}. Error: {e}")
|
|
68
|
+
|
|
69
|
+
if expected_hash and not verify_sha256(cached_file, expected_hash):
|
|
70
|
+
if os.path.exists(cached_file):
|
|
71
|
+
os.remove(cached_file)
|
|
72
|
+
raise ValueError(f"Checksum verification failed for {cached_file}. Cache cleared.")
|
|
73
|
+
|
|
74
|
+
def load_model(
|
|
75
|
+
variant: str,
|
|
76
|
+
device: str = "cpu",
|
|
77
|
+
force_download: bool = False,
|
|
78
|
+
architecture = None,
|
|
79
|
+
kwargs: dict = None,
|
|
80
|
+
backend: str = None,
|
|
81
|
+
sha256: str = None
|
|
82
|
+
) -> torch.nn.Module:
|
|
83
|
+
"""
|
|
84
|
+
Loads a model by its registry name OR directly from a remote HTTP(S) URL.
|
|
85
|
+
|
|
86
|
+
Args:
|
|
87
|
+
variant: Registry name (e.g. 'unet_cfd_v1') OR direct remote URL (e.g. 'https://.../model.pth').
|
|
88
|
+
device: The target device for execution (e.g., 'cpu', 'cuda', 'mps').
|
|
89
|
+
force_download: If True, deletes cached files and re-downloads weights.
|
|
90
|
+
architecture: Model architecture class (e.g., UNet) - required if loading PyTorch weights from a URL.
|
|
91
|
+
kwargs: Keyword arguments for instantiating the model class (used when loading from a URL).
|
|
92
|
+
backend: Backend target ('pytorch' or 'onnx'). Automatically inferred from URL if loading from a URL.
|
|
93
|
+
sha256: Optional SHA256 checksum to verify file integrity.
|
|
94
|
+
"""
|
|
95
|
+
is_url = variant.startswith("http://") or variant.startswith("https://")
|
|
96
|
+
|
|
97
|
+
if is_url:
|
|
98
|
+
url = variant
|
|
99
|
+
if backend is None:
|
|
100
|
+
backend = "onnx" if url.lower().endswith(".onnx") else "pytorch"
|
|
101
|
+
|
|
102
|
+
if backend == "pytorch" and architecture is None:
|
|
103
|
+
raise ValueError(
|
|
104
|
+
"You must specify the 'architecture' class (e.g., UNet, DeepCrack) "
|
|
105
|
+
"when loading a PyTorch model directly from a URL."
|
|
106
|
+
)
|
|
107
|
+
|
|
108
|
+
config = {
|
|
109
|
+
"architecture": architecture,
|
|
110
|
+
"kwargs": kwargs or {},
|
|
111
|
+
"backend": backend,
|
|
112
|
+
"url": url,
|
|
113
|
+
"sha256": sha256
|
|
114
|
+
}
|
|
115
|
+
else:
|
|
116
|
+
if variant not in MODEL_REGISTRY:
|
|
117
|
+
available_variants = list(MODEL_REGISTRY.keys())
|
|
118
|
+
raise ValueError(f"Unknown variant '{variant}'. Available: {available_variants}")
|
|
119
|
+
config = MODEL_REGISTRY[variant]
|
|
120
|
+
backend = config.get("backend", "pytorch")
|
|
121
|
+
|
|
122
|
+
# 1. Resolve caching paths
|
|
123
|
+
filename = os.path.basename(config["url"])
|
|
124
|
+
cached_file = os.path.join(get_cache_dir(), filename)
|
|
125
|
+
|
|
126
|
+
if force_download and os.path.exists(cached_file):
|
|
127
|
+
try:
|
|
128
|
+
os.remove(cached_file)
|
|
129
|
+
except OSError:
|
|
130
|
+
pass
|
|
131
|
+
|
|
132
|
+
# 2. Download weights if missing
|
|
133
|
+
if not os.path.exists(cached_file):
|
|
134
|
+
download_weight_file(config["url"], cached_file, config.get("sha256"))
|
|
135
|
+
|
|
136
|
+
# 3. Instantiate and load weights based on backend type
|
|
137
|
+
if backend == "pytorch":
|
|
138
|
+
arch_class = config["architecture"]
|
|
139
|
+
model = arch_class(**config.get("kwargs", {}))
|
|
140
|
+
|
|
141
|
+
# Load weights
|
|
142
|
+
state_dict = torch.load(cached_file, map_location=device)
|
|
143
|
+
model.load_state_dict(state_dict)
|
|
144
|
+
return model.to(device).eval()
|
|
145
|
+
|
|
146
|
+
elif backend == "onnx":
|
|
147
|
+
model = ONNXModelWrapper(cached_file, device=device)
|
|
148
|
+
return model
|
|
149
|
+
|
|
150
|
+
else:
|
|
151
|
+
raise ValueError(f"Unsupported backend: {backend}")
|
|
152
|
+
|
|
153
|
+
|
|
154
|
+
def list_models() -> list:
|
|
155
|
+
"""
|
|
156
|
+
Returns a list of all available pre-trained model variants in the registry.
|
|
157
|
+
"""
|
|
158
|
+
return list(MODEL_REGISTRY.keys())
|
|
159
|
+
|
|
160
|
+
|
|
161
|
+
def register_model(
|
|
162
|
+
name: str,
|
|
163
|
+
url: str,
|
|
164
|
+
architecture = None,
|
|
165
|
+
kwargs: dict = None,
|
|
166
|
+
backend: str = "pytorch",
|
|
167
|
+
sha256: str = None
|
|
168
|
+
):
|
|
169
|
+
"""
|
|
170
|
+
Dynamically registers a custom model variant at runtime.
|
|
171
|
+
This allows users to define custom remote models and use the standard load_model API.
|
|
172
|
+
|
|
173
|
+
Args:
|
|
174
|
+
name: Name/identifier of the variant.
|
|
175
|
+
url: Remote URL to download the model weights/file from.
|
|
176
|
+
architecture: The PyTorch class/type of the model (not required for ONNX).
|
|
177
|
+
kwargs: Keyword arguments for instantiating the model class.
|
|
178
|
+
backend: Backend framework, either 'pytorch' or 'onnx'.
|
|
179
|
+
sha256: Optional SHA256 checksum to verify the downloaded file.
|
|
180
|
+
"""
|
|
181
|
+
if backend not in ("pytorch", "onnx"):
|
|
182
|
+
raise ValueError("backend must be 'pytorch' or 'onnx'")
|
|
183
|
+
if backend == "pytorch" and architecture is None:
|
|
184
|
+
raise ValueError("architecture class must be provided for PyTorch backend")
|
|
185
|
+
|
|
186
|
+
MODEL_REGISTRY[name] = {
|
|
187
|
+
"architecture": architecture,
|
|
188
|
+
"kwargs": kwargs or {},
|
|
189
|
+
"backend": backend,
|
|
190
|
+
"url": url,
|
|
191
|
+
"sha256": sha256
|
|
192
|
+
}
|
|
@@ -0,0 +1,172 @@
|
|
|
1
|
+
import numpy as np
|
|
2
|
+
from typing import Tuple, Generator
|
|
3
|
+
|
|
4
|
+
"""
|
|
5
|
+
Version 1
|
|
6
|
+
"""
|
|
7
|
+
class PatchExtractor:
|
|
8
|
+
"""
|
|
9
|
+
Extracts the overlapping patches from a large image.
|
|
10
|
+
"""
|
|
11
|
+
def __init__(self, patch_size: Tuple[int, int], overlap_ratio: float = 0.2):
|
|
12
|
+
"""
|
|
13
|
+
Args:
|
|
14
|
+
patch_size (height, width): the size of the patch to be extracted.
|
|
15
|
+
overlap_ratio: float number between 0.0 and 0.99. the default value is 0.2,
|
|
16
|
+
meaning that the patches will overlap by 20% in both dimensions.
|
|
17
|
+
"""
|
|
18
|
+
|
|
19
|
+
if not (0.0 <= overlap_ratio < 1.0):
|
|
20
|
+
raise ValueError("overlap_ratio must be between 0.0 and 1.0")
|
|
21
|
+
|
|
22
|
+
self.patch_height, self.patch_width = patch_size
|
|
23
|
+
self.stride_height = int(self.patch_height * (1 - overlap_ratio))
|
|
24
|
+
self.stride_width = int(self.patch_width * (1 - overlap_ratio))
|
|
25
|
+
|
|
26
|
+
# ensure that stride is at least 1 to avoid infinite loops
|
|
27
|
+
self.stride_height = max(1, self.stride_height)
|
|
28
|
+
self.stride_width = max(1, self.stride_width)
|
|
29
|
+
|
|
30
|
+
def extract(self, image: np.ndarray) -> Generator[Tuple[np.ndarray, Tuple[int, int]], None, None]:
|
|
31
|
+
"""
|
|
32
|
+
Yields patches and their top-left (y, x) coordinates.
|
|
33
|
+
Handles edges by shifting the last patch to align with the image border.
|
|
34
|
+
"""
|
|
35
|
+
|
|
36
|
+
image_height, image_width = image.shape[:2]
|
|
37
|
+
seen_coordinates = set()
|
|
38
|
+
|
|
39
|
+
for y in range(0, image_height, self.stride_height):
|
|
40
|
+
for x in range(0, image_width, self.stride_width):
|
|
41
|
+
|
|
42
|
+
# shift the patch if it goes out of bounds
|
|
43
|
+
if y + self.patch_height > image_height:
|
|
44
|
+
y = image_height - self.patch_height
|
|
45
|
+
|
|
46
|
+
if x + self.patch_width > image_width:
|
|
47
|
+
x = image_width - self.patch_width
|
|
48
|
+
|
|
49
|
+
# ensure we don't yield the same patch multiple times
|
|
50
|
+
if (y, x) in seen_coordinates:
|
|
51
|
+
continue
|
|
52
|
+
seen_coordinates.add((y, x))
|
|
53
|
+
|
|
54
|
+
# yield the patch and its coordinates
|
|
55
|
+
yield image[y:y+self.patch_height, x:x+self.patch_width], (y, x)
|
|
56
|
+
|
|
57
|
+
# Patch Reconstruction
|
|
58
|
+
class PatchBlender:
|
|
59
|
+
"""
|
|
60
|
+
Reconstructs the full image from the overlapping patches using Gaussian Blending
|
|
61
|
+
to eliminate seam artifacts
|
|
62
|
+
"""
|
|
63
|
+
|
|
64
|
+
def __init__(self, output_shape: Tuple[int, int], patch_size: Tuple[int, int], num_classes: int = 1):
|
|
65
|
+
self.output_height, self.output_width = output_shape
|
|
66
|
+
self.patch_height, self.patch_width = patch_size
|
|
67
|
+
self.num_classes = num_classes
|
|
68
|
+
|
|
69
|
+
# accumulater for value and weights
|
|
70
|
+
# using float32 to avoid overflow during accumulation
|
|
71
|
+
self.accumulated_values = np.zeros((num_classes, self.output_height, self.output_width), dtype=np.float32)
|
|
72
|
+
self.accumilated_weights = np.zeros((num_classes, self.output_height, self.output_width), dtype=np.float32)
|
|
73
|
+
|
|
74
|
+
# pre-compute the 2d gaussian wieghts for a single patch
|
|
75
|
+
self.weight_map = self._create_gaussian_weight(self.patch_height, self.patch_width)
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
def _create_gaussian_weight(self, height: int, width: int, sigma_factor: float = 0.25) -> np.ndarray:
|
|
79
|
+
"""
|
|
80
|
+
Create a 2d Gaussian mask. Center center is 1.0, edge fade to 0.0
|
|
81
|
+
"""
|
|
82
|
+
x = np.linspace(-1, 1, width)
|
|
83
|
+
y = np.linspace(-1, 1, height)
|
|
84
|
+
|
|
85
|
+
X, Y = np.meshgrid(x, y)
|
|
86
|
+
|
|
87
|
+
# sigma_factor controls the drop-off. o.25 makes the edges fade to 0.0
|
|
88
|
+
weights = np.exp(-(X**2 + Y**2) / (2 * sigma_factor**2))
|
|
89
|
+
return weights.astype(np.float32)
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
def add_patch(self, patch_prediction: np.ndarray, coordinates: Tuple[int, int]):
|
|
93
|
+
"""
|
|
94
|
+
Adds a predicte patch to the accomultor.
|
|
95
|
+
"""
|
|
96
|
+
y, x = coordinates
|
|
97
|
+
c = patch_prediction.shape[0]
|
|
98
|
+
|
|
99
|
+
# expand the weight map to march number of classes if necessary
|
|
100
|
+
weights = np.expand_dims(self.weight_map, 0) if c == 1 else np.stack([self.weight_map] * c, axis=0)
|
|
101
|
+
|
|
102
|
+
self.accumulated_values[:, y:y+self.patch_height, x:x+self.patch_width] += patch_prediction * weights
|
|
103
|
+
self.accumilated_weights[:, y:y+self.patch_height, x:x+self.patch_width] += weights
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
def merge(self) -> np.ndarray:
|
|
107
|
+
"""
|
|
108
|
+
Return the final blended image of the shape (C, H, W)
|
|
109
|
+
"""
|
|
110
|
+
|
|
111
|
+
# avoid division by zero in areas with no patches
|
|
112
|
+
self.accumilated_weights[self.accumilated_weights == 0] = 1.0
|
|
113
|
+
|
|
114
|
+
result = self.accumulated_values / self.accumilated_weights
|
|
115
|
+
|
|
116
|
+
# if single class. squeeze the channel for convenience
|
|
117
|
+
if self.num_classes == 1:
|
|
118
|
+
return result.squeeze(0)
|
|
119
|
+
return result
|
|
120
|
+
|
|
121
|
+
|
|
122
|
+
"""
|
|
123
|
+
Version 2
|
|
124
|
+
"""
|
|
125
|
+
class SlidingWindowExtractor:
|
|
126
|
+
def __init__(self, patch_size: int, overlap_ratio: float = 0.2):
|
|
127
|
+
self.patch_size = patch_size
|
|
128
|
+
self.step_size = max(1, int(patch_size * (1 - overlap_ratio)))
|
|
129
|
+
|
|
130
|
+
def extract(self, image: np.ndarray) -> Generator[Tuple[np.ndarray, Tuple[int, int]], None, None]:
|
|
131
|
+
"""
|
|
132
|
+
yields patches and their (y, x) coordinates. automatically handles edges.
|
|
133
|
+
"""
|
|
134
|
+
height, width = image.shape[:2]
|
|
135
|
+
seen_coordinates = set()
|
|
136
|
+
|
|
137
|
+
for y in range(0, height, self.step_size):
|
|
138
|
+
for x in range(0, width, self.step_size):
|
|
139
|
+
# shift the patch if it goes out of bounds
|
|
140
|
+
if y + self.patch_size > height: y = height - self.patch_size
|
|
141
|
+
if x + self.patch_size > width: x = width - self.patch_size
|
|
142
|
+
|
|
143
|
+
if (y, x) in seen_coordinates: continue
|
|
144
|
+
seen_coordinates.add((y, x))
|
|
145
|
+
|
|
146
|
+
yield image[y:y+self.patch_size, x:x+self.patch_size], (y, x)
|
|
147
|
+
|
|
148
|
+
class CountMapBlender:
|
|
149
|
+
"""
|
|
150
|
+
Reconstructs the full image by averaging overlapping patches.
|
|
151
|
+
"""
|
|
152
|
+
def __init__(self, shape: Tuple[int, int]):
|
|
153
|
+
self.prediction_map = np.zeros(shape, dtype=np.float32)
|
|
154
|
+
self.count_map = np.zeros(shape, dtype=np.int32)
|
|
155
|
+
|
|
156
|
+
def add(self, patch: np.ndarray, coordinates: Tuple[int, int]):
|
|
157
|
+
"""
|
|
158
|
+
Adds a patch to the prediction map and updates the count map.
|
|
159
|
+
"""
|
|
160
|
+
y, x = coordinates
|
|
161
|
+
height, width = patch.shape
|
|
162
|
+
|
|
163
|
+
self.prediction_map[y:y+height, x:x+width] += patch
|
|
164
|
+
self.count_map[y:y+height, x:x+width] += 1
|
|
165
|
+
|
|
166
|
+
def merge(self) -> np.ndarray:
|
|
167
|
+
"""
|
|
168
|
+
Merges the prediction map with the count map to produce the final blended image.
|
|
169
|
+
"""
|
|
170
|
+
valid = self.count_map > 0
|
|
171
|
+
self.prediction_map[valid] /= self.count_map[valid]
|
|
172
|
+
return self.prediction_map
|
|
@@ -0,0 +1,89 @@
|
|
|
1
|
+
import cv2
|
|
2
|
+
import torch
|
|
3
|
+
import numpy as np
|
|
4
|
+
from pathlib import Path
|
|
5
|
+
from PIL import Image
|
|
6
|
+
|
|
7
|
+
from .preprocess import apply_lab_clahe, get_inference_transform
|
|
8
|
+
from .tta import tta_forward
|
|
9
|
+
from .patching import CountMapBlender, PatchExtractor, PatchBlender, SlidingWindowExtractor
|
|
10
|
+
from .models import load_model
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class CrackInferencePipeline:
|
|
14
|
+
"""
|
|
15
|
+
pipeline for running inference on large images.
|
|
16
|
+
"""
|
|
17
|
+
def __init__(self, model: torch.nn.Module, device: str = "cuda",
|
|
18
|
+
patch_size: int = 512, overlap_ratio: float = 0.2,
|
|
19
|
+
confidence_threhold: float = 0.5, use_tta: bool = False):
|
|
20
|
+
self.device = torch.device(device if torch.cuda.is_available() else "cpu")
|
|
21
|
+
self.model = model.to(self.device).eval()
|
|
22
|
+
self.patch_size = patch_size
|
|
23
|
+
self.overlap_ratio = overlap_ratio
|
|
24
|
+
self.confidence_threshold = confidence_threhold
|
|
25
|
+
self.use_tta = use_tta
|
|
26
|
+
|
|
27
|
+
self.extractor = SlidingWindowExtractor(self.patch_size, self.overlap_ratio)
|
|
28
|
+
self.transform = get_inference_transform()
|
|
29
|
+
|
|
30
|
+
@classmethod
|
|
31
|
+
def from_checkpoint(cls, model_class, checkpoint_path: str, **kwargs):
|
|
32
|
+
"""
|
|
33
|
+
Helper function to load model
|
|
34
|
+
"""
|
|
35
|
+
model = model_class(n_channels=3, n_classes=1) # Assuming binary segmentation
|
|
36
|
+
state_dict = torch.load(checkpoint_path, map_location="cpu")
|
|
37
|
+
model.load_state_dict(state_dict)
|
|
38
|
+
|
|
39
|
+
return cls(model, **kwargs)
|
|
40
|
+
|
|
41
|
+
@classmethod
|
|
42
|
+
def from_pretrained(cls, variant: str, device: str = "cuda", **kwargs):
|
|
43
|
+
"""
|
|
44
|
+
Helper function to load pretrained model from the model zoo.
|
|
45
|
+
"""
|
|
46
|
+
model = load_model(variant, device=device)
|
|
47
|
+
return cls(model, device=device, **kwargs)
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
@torch.no_grad()
|
|
51
|
+
def predict(self, image_path: str) -> dict:
|
|
52
|
+
"""
|
|
53
|
+
Runs full inference pipeline on a large image.
|
|
54
|
+
Returns a dictionary with the original image, probability map, and binary mask.
|
|
55
|
+
"""
|
|
56
|
+
# Load Image
|
|
57
|
+
original_image = np.array(Image.open(image_path).convert('RGB'))
|
|
58
|
+
height, width, _ = original_image.shape
|
|
59
|
+
|
|
60
|
+
# Preprocess (LAB-CLAHE)
|
|
61
|
+
preprocessed_image = apply_lab_clahe(original_image)
|
|
62
|
+
|
|
63
|
+
# Initialize Blender
|
|
64
|
+
blender = CountMapBlender(shape=(height, width))
|
|
65
|
+
|
|
66
|
+
# Sliding Window Inference
|
|
67
|
+
for patch_rgb, coordinates in self.extractor.extract(preprocessed_image):
|
|
68
|
+
# Transform to tensor
|
|
69
|
+
transformed = self.transform(image=patch_rgb)
|
|
70
|
+
patch_tensor = transformed["image"].to(self.device)
|
|
71
|
+
|
|
72
|
+
# Run Model
|
|
73
|
+
if self.use_tta:
|
|
74
|
+
pred_prob = tta_forward(self.model, patch_tensor)
|
|
75
|
+
else:
|
|
76
|
+
pred_prob = torch.sigmoid(self.model(patch_tensor.unsqueeze(0))).squeeze()
|
|
77
|
+
|
|
78
|
+
# Add to blender
|
|
79
|
+
blender.add(pred_prob.cpu().numpy(), coordinates)
|
|
80
|
+
|
|
81
|
+
# Merge and Threshold
|
|
82
|
+
confidence_map = blender.merge()
|
|
83
|
+
binary_mask = (confidence_map > self.confidence_threshold).astype(np.uint8) * 255
|
|
84
|
+
|
|
85
|
+
return {
|
|
86
|
+
"original_image": original_image,
|
|
87
|
+
"confidence_map": confidence_map,
|
|
88
|
+
"binary_mask": binary_mask
|
|
89
|
+
}
|
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
import cv2
|
|
2
|
+
import numpy as np
|
|
3
|
+
import albumentations as A
|
|
4
|
+
from albumentations.pytorch import ToTensorV2
|
|
5
|
+
|
|
6
|
+
def apply_lab_clahe(image: np.ndarray, clip_limit: float = 2.0) -> np.ndarray:
|
|
7
|
+
"""
|
|
8
|
+
Apply clahe to the L-channel of the LAB color space to enhance
|
|
9
|
+
local contrast without altering color balance.
|
|
10
|
+
"""
|
|
11
|
+
lab_image = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
|
|
12
|
+
l_channel, a_channel, b_channel = cv2.split(lab_image)
|
|
13
|
+
clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=(8, 8))
|
|
14
|
+
cl = clahe.apply(l_channel)
|
|
15
|
+
limg = cv2.merge((cl, a_channel, b_channel))
|
|
16
|
+
return cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)
|
|
17
|
+
|
|
18
|
+
def get_inference_transform():
|
|
19
|
+
"""
|
|
20
|
+
standard ImageNet normalization required by most pretrained models.
|
|
21
|
+
"""
|
|
22
|
+
return A.Compose([
|
|
23
|
+
A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
|
|
24
|
+
ToTensorV2(),
|
|
25
|
+
])
|
|
@@ -0,0 +1,40 @@
|
|
|
1
|
+
import torch
|
|
2
|
+
|
|
3
|
+
def tta_forward(model: torch.nn.Module, image_tensor: torch.Tensor) -> torch.Tensor:
|
|
4
|
+
"""
|
|
5
|
+
apply 4-way test-time augmentation and averages the sigmoid predictions.
|
|
6
|
+
expects image tensor to be shape (C, H, W). returns (H, W)
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
with torch.no_grad():
|
|
10
|
+
x = image_tensor.unsqueeze(0)
|
|
11
|
+
|
|
12
|
+
# original
|
|
13
|
+
original_prediction = torch.sigmoid(model(x))
|
|
14
|
+
|
|
15
|
+
# horizontal flip
|
|
16
|
+
horizontal_flip_prediction = torch.flip(torch.sigmoid(model(torch.flip(x, dims=[3]))), dims=[3])
|
|
17
|
+
|
|
18
|
+
# vertical flip
|
|
19
|
+
vertical_flip_prediction = torch.flip(torch.sigmoid(model(torch.flip(x, dims=[2]))), dims=[2])
|
|
20
|
+
|
|
21
|
+
# diagonal flip
|
|
22
|
+
# diagonal_flip_prediction = torch.flip(torch.sigmoid(model(torch.flip(x, dims=[2, 3]))), dims=[2, 3])
|
|
23
|
+
|
|
24
|
+
# average the predictions
|
|
25
|
+
# final_prediction = (original_prediction + horizontal_flip_prediction + vertical_flip_prediction + diagonal_flip_prediction) / 4
|
|
26
|
+
|
|
27
|
+
# rotate 90 degrees clockwise
|
|
28
|
+
rotated_90_prediction = torch.rot90(
|
|
29
|
+
torch.sigmoid(model(torch.rot90(x, k=1, dims=[2, 3]))),
|
|
30
|
+
k=-1, dims=[2,3]
|
|
31
|
+
)
|
|
32
|
+
|
|
33
|
+
# average predictions
|
|
34
|
+
averaged_prediction = (
|
|
35
|
+
original_prediction +
|
|
36
|
+
horizontal_flip_prediction +
|
|
37
|
+
vertical_flip_prediction +
|
|
38
|
+
rotated_90_prediction) / 4.0
|
|
39
|
+
|
|
40
|
+
return averaged_prediction.squeeze(0)
|
|
@@ -0,0 +1,152 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: findcrack
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: A deep learning crack detection package supporting U-Net and DeepCrack models with PyTorch and ONNX backends.
|
|
5
|
+
Author: StrikerEurika
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/StrikerEurika/findcrack
|
|
8
|
+
Project-URL: Repository, https://github.com/StrikerEurika/findcrack
|
|
9
|
+
Project-URL: Releases, https://github.com/StrikerEurika/findcrack/releases
|
|
10
|
+
Classifier: Programming Language :: Python :: 3
|
|
11
|
+
Classifier: Operating System :: OS Independent
|
|
12
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
13
|
+
Classifier: Topic :: Scientific/Engineering :: Image Processing
|
|
14
|
+
Requires-Python: >=3.11
|
|
15
|
+
Description-Content-Type: text/markdown
|
|
16
|
+
Requires-Dist: albumentations>=2.0.8
|
|
17
|
+
Requires-Dist: onnxruntime>=1.27.0
|
|
18
|
+
Requires-Dist: opencv-python>=4.13.0.92
|
|
19
|
+
Requires-Dist: torch>=2.6.0
|
|
20
|
+
Requires-Dist: torchaudio>=2.6.0
|
|
21
|
+
Requires-Dist: torchvision>=0.21.0
|
|
22
|
+
|
|
23
|
+
# findcrack
|
|
24
|
+
|
|
25
|
+
`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).
|
|
26
|
+
|
|
27
|
+
---
|
|
28
|
+
|
|
29
|
+
## Features
|
|
30
|
+
|
|
31
|
+
- **Pre-trained Model Zoo**: Fetch pre-trained model weights (e.g., `Seg_UNET_CFD_actual_v1`, `Seg_UNET_CFD_actual_v2`) dynamically on demand.
|
|
32
|
+
- **Unified Backend Engine**: Seamlessly executes either PyTorch (`.pth`/`.pt`) or ONNX (`.onnx`) models using the same standard interface.
|
|
33
|
+
- **Sliding-Window Inference**: Efficiently process ultra-high-resolution images by dividing them into overlapping patches.
|
|
34
|
+
- **Gaussian & Average Blending**: Reconstructs the full image from patches using overlapping Gaussian blending filters to eliminate edge-seam artifacts.
|
|
35
|
+
- **Test-Time Augmentation (TTA)**: Performs multi-way augmentations (original, horizontal flip, vertical flip, and rotations) to produce highly robust prediction masks.
|
|
36
|
+
- **Validation Metrics**: Compute standard segmentation metrics like IoU, Dice Coefficient, Precision, Recall, and Pixel Accuracy.
|
|
37
|
+
|
|
38
|
+
---
|
|
39
|
+
|
|
40
|
+
## Installation
|
|
41
|
+
|
|
42
|
+
You can install `findcrack` directly from source or via PyPI (once published):
|
|
43
|
+
|
|
44
|
+
```bash
|
|
45
|
+
# Install via pip
|
|
46
|
+
pip install findcrack
|
|
47
|
+
|
|
48
|
+
# Or using uv
|
|
49
|
+
uv add findcrack
|
|
50
|
+
```
|
|
51
|
+
|
|
52
|
+
---
|
|
53
|
+
|
|
54
|
+
## Quickstart
|
|
55
|
+
|
|
56
|
+
Here is how to load a pre-trained model and run crack detection on a large image:
|
|
57
|
+
|
|
58
|
+
```python
|
|
59
|
+
import cv2
|
|
60
|
+
from findcrack import CrackInferencePipeline, load_model
|
|
61
|
+
|
|
62
|
+
# 1. Load a pre-trained model from the official registry (or use your own URL)
|
|
63
|
+
# The weights are downloaded dynamically from GitHub releases on first use.
|
|
64
|
+
model = load_model("Seg_UNET_CFD_actual_v1", device="cuda")
|
|
65
|
+
|
|
66
|
+
# 2. Setup the inference pipeline
|
|
67
|
+
pipeline = CrackInferencePipeline(
|
|
68
|
+
model=model,
|
|
69
|
+
device="cuda",
|
|
70
|
+
patch_size=512,
|
|
71
|
+
overlap_ratio=0.2,
|
|
72
|
+
confidence_threhold=0.5,
|
|
73
|
+
use_tta=True # Enables multi-way Test-Time Augmentation
|
|
74
|
+
)
|
|
75
|
+
|
|
76
|
+
# 3. Perform inference
|
|
77
|
+
results = pipeline.predict("path/to/high_res_concrete.jpg")
|
|
78
|
+
|
|
79
|
+
# The results dictionary contains:
|
|
80
|
+
# - results["original_image"]: Original RGB image (numpy array)
|
|
81
|
+
# - results["confidence_map"]: Float probability map [0.0 - 1.0]
|
|
82
|
+
# - results["binary_mask"]: Binary segmentation mask [0 or 255]
|
|
83
|
+
|
|
84
|
+
# Save the output mask
|
|
85
|
+
cv2.imwrite("detected_cracks.png", results["binary_mask"])
|
|
86
|
+
```
|
|
87
|
+
|
|
88
|
+
---
|
|
89
|
+
|
|
90
|
+
## API Reference
|
|
91
|
+
|
|
92
|
+
### Model Loading & Caching
|
|
93
|
+
|
|
94
|
+
#### `load_model(variant: str, device: str = "cpu", force_download: bool = False, architecture = None, **kwargs)`
|
|
95
|
+
Loads a model variant from the local registry or directly from a remote HTTP(S) URL.
|
|
96
|
+
|
|
97
|
+
- **Parameters**:
|
|
98
|
+
- `variant`: The name of a registered variant (e.g., `"Seg_UNET_CFD_actual_v1"`) or a direct HTTP(S) URL to a weights file.
|
|
99
|
+
- `device`: Target execution device (`"cpu"`, `"cuda"`, or `"mps"`).
|
|
100
|
+
- `force_download`: If `True`, re-downloads weights even if cached locally.
|
|
101
|
+
- `architecture`: PyTorch architecture class (e.g., `UNet`, `DeepCrack`) - required only if loading a raw `.pth`/`.pt` file from a custom URL.
|
|
102
|
+
|
|
103
|
+
```python
|
|
104
|
+
from findcrack import load_model, UNet
|
|
105
|
+
|
|
106
|
+
# Load custom model weights directly from an external URL
|
|
107
|
+
model = load_model(
|
|
108
|
+
variant="https://my-domain.com/custom_unet.pth",
|
|
109
|
+
architecture=UNet,
|
|
110
|
+
device="cuda"
|
|
111
|
+
)
|
|
112
|
+
```
|
|
113
|
+
|
|
114
|
+
#### `list_models()`
|
|
115
|
+
Returns a list of all pre-trained models available in the built-in registry.
|
|
116
|
+
|
|
117
|
+
#### `register_model(name: str, url: str, architecture = None, kwargs: dict = None, backend: str = "pytorch")`
|
|
118
|
+
Registers a custom variant dynamically at runtime.
|
|
119
|
+
|
|
120
|
+
---
|
|
121
|
+
|
|
122
|
+
### Pipeline Configuration
|
|
123
|
+
|
|
124
|
+
#### `CrackInferencePipeline(model, device: str = "cuda", patch_size: int = 512, overlap_ratio: float = 0.2, confidence_threhold: float = 0.5, use_tta: bool = False)`
|
|
125
|
+
Handles sliding window preprocessing, execution, TTA, and patching reconstruction.
|
|
126
|
+
|
|
127
|
+
---
|
|
128
|
+
|
|
129
|
+
## Directory Structure
|
|
130
|
+
|
|
131
|
+
```text
|
|
132
|
+
src/
|
|
133
|
+
└── findcrack/
|
|
134
|
+
├── __init__.py # Main API endpoints (load_model, CrackInferencePipeline, etc.)
|
|
135
|
+
├── metrics.py # Segmentation evaluation metrics (IoU, Dice, etc.)
|
|
136
|
+
├── patching.py # Sliding window extraction and blend reconstruction
|
|
137
|
+
├── pipeline.py # Crack Inference Pipeline wrapper
|
|
138
|
+
├── preprocess.py # Color-space CLAHE contrast enhancement & transforms
|
|
139
|
+
├── tta.py # Test-Time Augmentation forward pass routines
|
|
140
|
+
└── models/
|
|
141
|
+
├── __init__.py # Model module exports
|
|
142
|
+
├── unet.py # U-Net model definition
|
|
143
|
+
├── deepcrack.py # DeepCrack model definition
|
|
144
|
+
├── onnx_wrapper.py # Wrapper for running ONNX models as nn.Modules
|
|
145
|
+
└── zoo.py # Remote weight registry and cached loaders
|
|
146
|
+
```
|
|
147
|
+
|
|
148
|
+
---
|
|
149
|
+
|
|
150
|
+
## License
|
|
151
|
+
|
|
152
|
+
This project is licensed under the MIT License.
|
|
@@ -0,0 +1,18 @@
|
|
|
1
|
+
README.md
|
|
2
|
+
pyproject.toml
|
|
3
|
+
src/findcrack/__init__.py
|
|
4
|
+
src/findcrack/metrics.py
|
|
5
|
+
src/findcrack/patching.py
|
|
6
|
+
src/findcrack/pipeline.py
|
|
7
|
+
src/findcrack/preprocess.py
|
|
8
|
+
src/findcrack/tta.py
|
|
9
|
+
src/findcrack.egg-info/PKG-INFO
|
|
10
|
+
src/findcrack.egg-info/SOURCES.txt
|
|
11
|
+
src/findcrack.egg-info/dependency_links.txt
|
|
12
|
+
src/findcrack.egg-info/requires.txt
|
|
13
|
+
src/findcrack.egg-info/top_level.txt
|
|
14
|
+
src/findcrack/models/__init__.py
|
|
15
|
+
src/findcrack/models/deepcrack.py
|
|
16
|
+
src/findcrack/models/onnx_wrapper.py
|
|
17
|
+
src/findcrack/models/unet.py
|
|
18
|
+
src/findcrack/models/zoo.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
findcrack
|