Coverage for src/crackmnist_eccv/run_experiment.py: 0%
133 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-10-19 08:00 +0000
1# SPDX-FileCopyrightText: 2026 German Aerospace Center (DLR)
2# SPDX-License-Identifier: GPL-3.0-or-later
4"""Run CrackMNIST pollution experiments across seeds, intensities, and splits.
6Trains a UNet model under controlled input pollution and evaluates average
7Dice loss on train/val/test splits across multiple random seeds. Results are
8written incrementally to a CSV file named after the experiment configuration.
9"""
11import argparse
12import ast
13import random
14import time
16from crackmnist import CrackMNIST
17from crackmnist.loss import DiceLoss
18from crackmnist.nets import UNet
19from crackmnist.transforms import InputNormalization, NumpyToTensor
20import numpy as np
21import pandas as pd
22import torch
23from torch import optim
24from torch.utils import data
25from torchvision import transforms
27from src.crackmnist_eccv.utility import apply_per_image_pollution, enlarge_crack_tip
30##############################################################################
31# FIXED HYPERPARAMETERS
32##############################################################################
33NUM_EPOCHS = 10
34BATCH_SIZE = 32
35LR = 1e-4
37##############################################################################
38# PARAM RANGES
39##############################################################################
40TEST_PERC_RANGE = [0.0, 1.0]
41NUM_SEEDS = 10
42SIZE = "S"
45##############################################################################
46# COMMAND-LINE ARGUMENTS
47##############################################################################
48def parse_python_literal(text: str, expected_type: type) -> list:
49 """Convert a string like '[0.1,0.2]' or '[(0.1,0.2)]' to a Python object.
51 Args:
52 text: String representation of a Python literal.
53 expected_type: The expected Python type (e.g. list).
55 Returns:
56 Parsed Python object of the expected type.
58 Raises:
59 ArgumentTypeError: If parsing fails or type does not match.
60 TypeError: If the parsed value is not an instance of expected_type.
61 """
62 try:
63 value = ast.literal_eval(text)
64 if not isinstance(value, expected_type):
65 raise TypeError
66 return value
67 except Exception as err:
68 raise argparse.ArgumentTypeError(
69 f"Value '{text}' is not a valid {expected_type.__name__}"
70 ) from err
73parser = argparse.ArgumentParser(
74 description="Run CrackMNIST experiments with flexible settings."
75)
77parser.add_argument(
78 "pollution_type",
79 nargs="?",
80 default="blob",
81 help="gaussian | salt_pepper | gaussian_blur | blob",
82)
83parser.add_argument(
84 "intensity_range",
85 nargs="?",
86 default="[(0.1,0.2),(0.3,0.4),(0.5,0.6)]",
87 help='Python literal, e.g. "[0.1,0.3,0.5]" or "[(0.1,0.2)]"',
88)
89parser.add_argument(
90 "--pixels",
91 type=int,
92 default=64,
93 help="Image size fed into CrackMNIST (default: 64)",
94)
95parser.add_argument(
96 "--train_range",
97 type=lambda s: parse_python_literal(s, list),
98 default=[0.05, 0.1, 0.15, 0.2],
99 help='Training-pollution percentages, e.g. "[0.0,0.25,0.5]"',
100)
101parser.add_argument(
102 "--radius",
103 type=int,
104 default=8,
105 help=("Blob-pollution radius (only used when pollution_type == 'blob'; default: 8)"),
106)
107args = parser.parse_args()
109##############################################################################
110# SET VARIABLES FROM ARGS
111##############################################################################
112SELECTED_POLLUTION_TYPE = args.pollution_type
113INTENSITY_RANGE = parse_python_literal(args.intensity_range, list)
114PIXELS = args.pixels
115TRAIN_PERC_RANGE = args.train_range
116RADIUS = args.radius
118NAMING_CONSTANT = f"{PIXELS}_F9_{SELECTED_POLLUTION_TYPE}_radius{RADIUS}"
121##############################################################################
122# UTILITY: FIX RANDOM SEEDS
123##############################################################################
124def fix_random_seeds(seed: int) -> None:
125 """Fix all relevant random seeds for reproducibility.
127 Args:
128 seed: Integer seed value.
129 """
130 random.seed(seed)
131 np.random.seed(seed)
132 torch.manual_seed(seed)
133 torch.use_deterministic_algorithms(True)
136##############################################################################
137# DEVICE
138##############################################################################
139device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
140print(f"Using device: {device}")
142##############################################################################
143# LOAD DATASETS
144##############################################################################
145data_transform = transforms.Compose([
146 NumpyToTensor(),
147 InputNormalization(),
148])
149target_transform = transforms.Compose([
150 NumpyToTensor(),
151])
153train_dataset_full = CrackMNIST(
154 split="train",
155 pixels=PIXELS,
156 size=SIZE,
157 transform=data_transform,
158 target_transform=target_transform,
159)
161val_dataset_full = CrackMNIST(
162 split="val",
163 pixels=PIXELS,
164 size=SIZE,
165 transform=data_transform,
166 target_transform=target_transform,
167)
169test_dataset_full = CrackMNIST(
170 split="test",
171 pixels=PIXELS,
172 size=SIZE,
173 transform=data_transform,
174 target_transform=target_transform,
175)
177##############################################################################
178# ENLARGE THE ANNOTATED TIP REGION TO 3x3
179##############################################################################
180train_dataset_full.masks = np.array([
181 enlarge_crack_tip(mask) for mask in train_dataset_full.masks
182])
183val_dataset_full.masks = np.array([
184 enlarge_crack_tip(mask) for mask in val_dataset_full.masks
185])
186test_dataset_full.masks = np.array([
187 enlarge_crack_tip(mask) for mask in test_dataset_full.masks
188])
190print("Dataset loaded")
191train_loader_full = data.DataLoader(
192 train_dataset_full, batch_size=BATCH_SIZE, shuffle=False, num_workers=4
193)
194val_loader_full = data.DataLoader(
195 val_dataset_full, batch_size=BATCH_SIZE, shuffle=False, num_workers=4
196)
197test_loader_full = data.DataLoader(
198 test_dataset_full, batch_size=BATCH_SIZE, shuffle=False, num_workers=4
199)
202##############################################################################
203# HELPER: CONFIG & EVALUATION FUNCTION
204##############################################################################
205def get_pollution_config(
206 pollution_type: str, intensity: float | tuple[float, float]
207) -> dict:
208 """Return keyword arguments for apply_per_image_pollution.
210 Args:
211 pollution_type: One of 'blob', 'gaussian', 'salt_pepper',
212 'gaussian_blur'.
213 intensity: Float or (min, max) tuple depending on pollution type.
215 Returns:
216 Dict with 'pollution_type' and 'pollution_kwargs' keys.
218 Raises:
219 ValueError: If the pollution type is unsupported or intensity format
220 is wrong.
221 """
222 if pollution_type == "blob":
223 if not isinstance(intensity, tuple) or len(intensity) != 2:
224 raise ValueError("Blob pollution requires intensity as a (min, max) tuple")
225 return {
226 "pollution_type": pollution_type,
227 "pollution_kwargs": {
228 "radius": RADIUS,
229 "target_min_oar": intensity[0],
230 "target_max_oar": intensity[1],
231 },
232 }
233 elif pollution_type in ["gaussian", "salt_pepper", "gaussian_blur"]:
234 if not isinstance(intensity, (float, int)):
235 raise ValueError(f"{pollution_type} requires intensity as a float or int")
236 return {
237 "pollution_type": pollution_type,
238 "pollution_kwargs": {"intensity": intensity},
239 }
240 else:
241 raise ValueError(f"Unsupported pollution type: {pollution_type}")
244def evaluate(
245 model: torch.nn.Module,
246 loader: data.DataLoader,
247 pollution_ratio: float,
248 pollution_type: str,
249 pollution_kwargs: dict,
250) -> float:
251 """Run model in eval mode and return average DiceLoss over the loader.
253 Args:
254 model: Trained PyTorch model.
255 loader: DataLoader to evaluate on.
256 pollution_ratio: Fraction of images to pollute.
257 pollution_type: Type of pollution to apply.
258 pollution_kwargs: Additional keyword arguments for the pollution
259 function.
261 Returns:
262 Average DiceLoss as a float.
263 """
264 model.eval()
265 criterion = DiceLoss()
266 total_loss, total_count = 0.0, 0
268 with torch.no_grad():
269 for inputs, targets in loader:
270 inputs = apply_per_image_pollution(
271 inputs,
272 pollution_ratio=pollution_ratio,
273 pollution_type=pollution_type,
274 **pollution_kwargs,
275 )
276 inputs, targets = inputs.to(device), targets.to(device)
277 outputs = model(inputs)
278 loss = criterion(outputs, targets.unsqueeze(1).float() / 255.0)
279 batch_size = inputs.size(0)
280 total_loss += loss.item() * batch_size
281 total_count += batch_size
283 return total_loss / total_count if total_count else 0.0
286##############################################################################
287# TRAIN & FINAL-EVAL
288##############################################################################
289def train_and_evaluate(
290 train_perc: float,
291 test_perc: float,
292 intensity: float | tuple[float, float],
293 seed: int,
294) -> tuple[float, float, float]:
295 """Train a UNet and evaluate on val and test sets.
297 Args:
298 train_perc: Fraction of training images to pollute.
299 test_perc: Fraction of test images to pollute during evaluation.
300 intensity: Pollution intensity (float or tuple for blob).
301 seed: Random seed for reproducibility.
303 Returns:
304 Tuple of (final_val_loss, final_test_loss, final_train_loss).
305 """
306 fix_random_seeds(seed)
308 model = UNet(in_ch=2, out_ch=1, init_features=128, dropout_prob=0.2).to(device)
309 criterion = DiceLoss()
310 optimizer = optim.Adam(model.parameters(), lr=LR)
312 config = get_pollution_config(SELECTED_POLLUTION_TYPE, intensity)
313 pollution_type = config["pollution_type"]
314 pollution_kwargs = config["pollution_kwargs"]
316 model.train()
317 for _epoch in range(NUM_EPOCHS):
318 for inputs, targets in train_loader_full:
319 inputs = apply_per_image_pollution(
320 inputs,
321 pollution_ratio=train_perc,
322 pollution_type=pollution_type,
323 **pollution_kwargs,
324 )
325 inputs, targets = inputs.to(device), targets.to(device)
327 optimizer.zero_grad()
328 outputs = model(inputs)
329 loss = criterion(outputs, targets.unsqueeze(1).float() / 255.0)
330 loss.backward()
331 optimizer.step()
333 final_val_loss = evaluate(
334 model,
335 val_loader_full,
336 pollution_ratio=train_perc,
337 pollution_type=pollution_type,
338 pollution_kwargs=pollution_kwargs,
339 )
340 final_test_loss = evaluate(
341 model,
342 test_loader_full,
343 pollution_ratio=test_perc,
344 pollution_type=pollution_type,
345 pollution_kwargs=pollution_kwargs,
346 )
347 final_train_loss = evaluate(
348 model,
349 train_loader_full,
350 pollution_ratio=train_perc,
351 pollution_type=pollution_type,
352 pollution_kwargs=pollution_kwargs,
353 )
355 return final_val_loss, final_test_loss, final_train_loss
358##############################################################################
359# MAIN LOOP
360##############################################################################
361results = []
363for train_perc in TRAIN_PERC_RANGE:
364 for test_perc in TEST_PERC_RANGE:
365 for intensity in INTENSITY_RANGE:
366 val_losses_seeds = []
367 test_losses_seeds = []
368 train_losses_seeds = []
369 print(f"Processing: {train_perc}, {test_perc}, {intensity}")
371 for seed in range(NUM_SEEDS):
372 print("\t seed:", seed)
373 start = time.time()
375 val_loss, test_loss, train_loss = train_and_evaluate(
376 train_perc, test_perc, intensity=intensity, seed=seed
377 )
378 val_losses_seeds.append(val_loss)
379 test_losses_seeds.append(test_loss)
380 train_losses_seeds.append(train_loss)
381 print("\t elapsed time:", time.time() - start)
383 avg_val_loss = np.mean(val_losses_seeds)
384 std_val_loss = np.std(val_losses_seeds)
385 avg_test_loss = np.mean(test_losses_seeds)
386 std_test_loss = np.std(test_losses_seeds)
387 avg_train_loss = np.mean(train_losses_seeds)
388 std_train_loss = np.std(train_losses_seeds)
390 results.append({
391 "train_perc": train_perc,
392 "test_perc": test_perc,
393 "intensity": intensity,
394 "avg_test_loss": avg_test_loss,
395 "std_test_loss": std_test_loss,
396 "avg_val_loss": avg_val_loss,
397 "std_val_loss": std_val_loss,
398 "avg_train_loss": avg_train_loss,
399 "std_train_loss": std_train_loss,
400 "test_losses_seeds": test_losses_seeds,
401 })
403 df = pd.DataFrame(results)
404 df.to_csv(f"{NAMING_CONSTANT}.csv", index=False)
406print("DONE!")