structrestore-3d 0.1.0a1__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.
Files changed (24) hide show
  1. structrestore_3d-0.1.0a1/LICENSE +21 -0
  2. structrestore_3d-0.1.0a1/PKG-INFO +80 -0
  3. structrestore_3d-0.1.0a1/README.md +55 -0
  4. structrestore_3d-0.1.0a1/pyproject.toml +39 -0
  5. structrestore_3d-0.1.0a1/setup.cfg +4 -0
  6. structrestore_3d-0.1.0a1/src/structrestore3d/__init__.py +3 -0
  7. structrestore_3d-0.1.0a1/src/structrestore3d/__main__.py +4 -0
  8. structrestore_3d-0.1.0a1/src/structrestore3d/audit_masking.py +105 -0
  9. structrestore_3d-0.1.0a1/src/structrestore3d/cli.py +64 -0
  10. structrestore_3d-0.1.0a1/src/structrestore3d/compare_v0_v1.py +213 -0
  11. structrestore_3d-0.1.0a1/src/structrestore3d/diagnose_v1_inference.py +260 -0
  12. structrestore_3d-0.1.0a1/src/structrestore3d/infer_baseline.py +88 -0
  13. structrestore_3d-0.1.0a1/src/structrestore3d/step29_registration_probe.py +244 -0
  14. structrestore_3d-0.1.0a1/src/structrestore3d/structrestore_25d.py +150 -0
  15. structrestore_3d-0.1.0a1/src/structrestore3d/train_baseline.py +75 -0
  16. structrestore_3d-0.1.0a1/src/structrestore3d/train_fullres.py +78 -0
  17. structrestore_3d-0.1.0a1/src/structrestore3d/train_v1_audited.py +130 -0
  18. structrestore_3d-0.1.0a1/src/structrestore3d/validate_25d.py +120 -0
  19. structrestore_3d-0.1.0a1/src/structrestore_3d.egg-info/PKG-INFO +80 -0
  20. structrestore_3d-0.1.0a1/src/structrestore_3d.egg-info/SOURCES.txt +22 -0
  21. structrestore_3d-0.1.0a1/src/structrestore_3d.egg-info/dependency_links.txt +1 -0
  22. structrestore_3d-0.1.0a1/src/structrestore_3d.egg-info/entry_points.txt +2 -0
  23. structrestore_3d-0.1.0a1/src/structrestore_3d.egg-info/requires.txt +6 -0
  24. structrestore_3d-0.1.0a1/src/structrestore_3d.egg-info/top_level.txt +1 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Krishna Kant Gupta
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,80 @@
1
+ Metadata-Version: 2.4
2
+ Name: structrestore-3d
3
+ Version: 0.1.0a1
4
+ Summary: Auditable self-supervised 2.5D FIB-SEM restoration research prototype
5
+ Author-email: Krishna Kant Gupta <krishnakant.gupta@bharatividyapeeth.edu>
6
+ Project-URL: Homepage, https://github.com/krishnakantgupta-ai/StructRestore-3D
7
+ Project-URL: Repository, https://github.com/krishnakantgupta-ai/StructRestore-3D
8
+ Project-URL: Issues, https://github.com/krishnakantgupta-ai/StructRestore-3D/issues
9
+ Keywords: FIB-SEM,electron microscopy,self-supervised learning,denoising
10
+ Classifier: Development Status :: 3 - Alpha
11
+ Classifier: Intended Audience :: Science/Research
12
+ Classifier: Programming Language :: Python :: 3
13
+ Classifier: Programming Language :: Python :: 3.12
14
+ Classifier: Topic :: Scientific/Engineering :: Image Processing
15
+ Requires-Python: <3.13,>=3.10
16
+ Description-Content-Type: text/markdown
17
+ License-File: LICENSE
18
+ Requires-Dist: numpy<3,>=1.24
19
+ Requires-Dist: scipy<2,>=1.10
20
+ Requires-Dist: tifffile>=2023.7.10
21
+ Requires-Dist: scikit-image<1,>=0.21
22
+ Requires-Dist: matplotlib<4,>=3.7
23
+ Requires-Dist: torch<3,>=2.2.2
24
+ Dynamic: license-file
25
+
26
+ # StructRestore-3D
27
+
28
+ Self-Supervised Deep Learning for FIB-SEM Microscopy Restoration
29
+
30
+ ## Overview
31
+
32
+ StructRestore-3D is a research prototype for investigating
33
+ self-supervised restoration of three-dimensional FIB-SEM
34
+ microscopy images while preserving cellular ultrastructure.
35
+
36
+ ## Implemented Models
37
+
38
+ 1. Single-slice convolutional neural network
39
+ 2. Plain 2.5D convolutional neural network
40
+ 3. Reliability-gated 2.5D convolutional neural network
41
+
42
+ ## Dataset
43
+
44
+ EMPIAR-10311: HeLa FIB-SEM microscopy dataset.
45
+
46
+ DOI: https://doi.org/10.6019/EMPIAR-10311
47
+
48
+ ## Features
49
+
50
+ - Self-supervised model training
51
+ - Masked-pixel reconstruction
52
+ - Direct target-leakage auditing
53
+ - Patch-quality diagnostics
54
+ - Multi-slice inference
55
+ - Quantitative benchmarking using PSNR and SSIM
56
+ - Neighboring-section alignment analysis
57
+
58
+ ## Preliminary Results
59
+
60
+ Six models were evaluated against noisy-input and conventional
61
+ filtering baselines. The present models have not demonstrated
62
+ superior denoising performance. Fine-structure preservation
63
+ remains an important research challenge.
64
+
65
+ ## Requirements
66
+
67
+ Python 3.12, PyTorch, NumPy, SciPy, tifffile,
68
+ scikit-image and Matplotlib.
69
+
70
+ ## Author
71
+
72
+ Dr. Krishna Kant Gupta
73
+
74
+ Rajiv Gandhi Institute of Information Technology and Biotechnology
75
+
76
+ Bharati Vidyapeeth (Deemed to be University), Pune, India
77
+
78
+ ## Research Status
79
+
80
+ Preliminary proof of concept. Further validation is ongoing.
@@ -0,0 +1,55 @@
1
+ # StructRestore-3D
2
+
3
+ Self-Supervised Deep Learning for FIB-SEM Microscopy Restoration
4
+
5
+ ## Overview
6
+
7
+ StructRestore-3D is a research prototype for investigating
8
+ self-supervised restoration of three-dimensional FIB-SEM
9
+ microscopy images while preserving cellular ultrastructure.
10
+
11
+ ## Implemented Models
12
+
13
+ 1. Single-slice convolutional neural network
14
+ 2. Plain 2.5D convolutional neural network
15
+ 3. Reliability-gated 2.5D convolutional neural network
16
+
17
+ ## Dataset
18
+
19
+ EMPIAR-10311: HeLa FIB-SEM microscopy dataset.
20
+
21
+ DOI: https://doi.org/10.6019/EMPIAR-10311
22
+
23
+ ## Features
24
+
25
+ - Self-supervised model training
26
+ - Masked-pixel reconstruction
27
+ - Direct target-leakage auditing
28
+ - Patch-quality diagnostics
29
+ - Multi-slice inference
30
+ - Quantitative benchmarking using PSNR and SSIM
31
+ - Neighboring-section alignment analysis
32
+
33
+ ## Preliminary Results
34
+
35
+ Six models were evaluated against noisy-input and conventional
36
+ filtering baselines. The present models have not demonstrated
37
+ superior denoising performance. Fine-structure preservation
38
+ remains an important research challenge.
39
+
40
+ ## Requirements
41
+
42
+ Python 3.12, PyTorch, NumPy, SciPy, tifffile,
43
+ scikit-image and Matplotlib.
44
+
45
+ ## Author
46
+
47
+ Dr. Krishna Kant Gupta
48
+
49
+ Rajiv Gandhi Institute of Information Technology and Biotechnology
50
+
51
+ Bharati Vidyapeeth (Deemed to be University), Pune, India
52
+
53
+ ## Research Status
54
+
55
+ Preliminary proof of concept. Further validation is ongoing.
@@ -0,0 +1,39 @@
1
+ [build-system]
2
+ requires = ["setuptools>=77", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "structrestore-3d"
7
+ version = "0.1.0a1"
8
+ description = "Auditable self-supervised 2.5D FIB-SEM restoration research prototype"
9
+ readme = "README.md"
10
+ requires-python = ">=3.10,<3.13"
11
+ authors = [{name = "Krishna Kant Gupta", email = "krishnakant.gupta@bharatividyapeeth.edu"}]
12
+ keywords = ["FIB-SEM", "electron microscopy", "self-supervised learning", "denoising"]
13
+ classifiers = [
14
+ "Development Status :: 3 - Alpha",
15
+ "Intended Audience :: Science/Research",
16
+ "Programming Language :: Python :: 3",
17
+ "Programming Language :: Python :: 3.12",
18
+ "Topic :: Scientific/Engineering :: Image Processing",
19
+ ]
20
+ dependencies = [
21
+ "numpy>=1.24,<3",
22
+ "scipy>=1.10,<2",
23
+ "tifffile>=2023.7.10",
24
+ "scikit-image>=0.21,<1",
25
+ "matplotlib>=3.7,<4",
26
+ "torch>=2.2.2,<3",
27
+ ]
28
+
29
+ [project.urls]
30
+ Homepage = "https://github.com/krishnakantgupta-ai/StructRestore-3D"
31
+ Repository = "https://github.com/krishnakantgupta-ai/StructRestore-3D"
32
+ Issues = "https://github.com/krishnakantgupta-ai/StructRestore-3D/issues"
33
+
34
+ [project.scripts]
35
+ structrestore3d = "structrestore3d.cli:main"
36
+
37
+ [tool.setuptools.packages.find]
38
+ where = ["src"]
39
+ include = ["structrestore3d*"]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,3 @@
1
+ """StructRestore-3D research prototype; methods are not yet validated denoisers."""
2
+
3
+ __version__ = "0.1.0a1"
@@ -0,0 +1,4 @@
1
+ from .cli import main
2
+
3
+ if __name__ == "__main__":
4
+ raise SystemExit(main())
@@ -0,0 +1,105 @@
1
+ #!/usr/bin/env python3
2
+ """Read-only audit of StructRestore-3D masking, split, and neighbor dependence.
3
+ Run from project root. Does not train or overwrite checkpoints.
4
+ """
5
+ import argparse, csv, json, sys, random
6
+ from pathlib import Path
7
+ import numpy as np
8
+ import tifffile
9
+ import torch
10
+
11
+ def main():
12
+ ap=argparse.ArgumentParser()
13
+ ap.add_argument('--input',default='data/raw')
14
+ ap.add_argument('--script',default='scripts/structrestore_25d.py')
15
+ ap.add_argument('--checkpoint',default='results/25d_gated/model.pt')
16
+ ap.add_argument('--patch',type=int,default=128)
17
+ ap.add_argument('--samples',type=int,default=100)
18
+ ap.add_argument('--out',default='results/masking_audit')
19
+ a=ap.parse_args()
20
+ torch.set_num_threads(2)
21
+ random.seed(42); np.random.seed(42); torch.manual_seed(42)
22
+ out=Path(a.out);out.mkdir(parents=True,exist_ok=True)
23
+ files=sorted(Path(a.input).glob('slice_*.tif'))
24
+ assert len(files)>=7, 'Expected at least seven TIFF slices'
25
+ train,test=files[:-4],files[-4:]
26
+ print('Training slices:',[f.name for f in train])
27
+ print('Held-out slices:',[f.name for f in test])
28
+ print('Filename overlap:',bool(set(train)&set(test)))
29
+ print('Adjacent train/test boundary:',train[-1].name,'/',test[0].name)
30
+ print('CAUTION: Adjacent sections are not biologically independent.')
31
+ print('CAUTION: Noisy neighboring sections can reveal central pixel structure; not a strict blind-spot model.')
32
+ import importlib.util
33
+ spec=importlib.util.spec_from_file_location('audit_model',a.script)
34
+ module=importlib.util.module_from_spec(spec);spec.loader.exec_module(module)
35
+ ck=torch.load(a.checkpoint,map_location='cpu',weights_only=True)
36
+ model=module.Net25D(ck['mode'],ck['width'])
37
+ model.load_state_dict(ck['state_dict']);model.eval()
38
+ print('Checkpoint mode:',ck['mode'],'width:',ck['width'])
39
+ # Check center masking: selected target pixels must differ from corrupted center in most cases.
40
+ ims=[tifffile.memmap(str(f)) for f in train]
41
+ rows=[]
42
+ with torch.inference_mode():
43
+ for i in range(a.samples):
44
+ z=random.randrange(1,len(ims)-1)
45
+ h,w=ims[z].shape
46
+ y=random.randrange(h-a.patch+1);x=random.randrange(w-a.patch+1)
47
+ arr=np.stack([np.asarray(ims[j][y:y+a.patch,x:x+a.patch],dtype=np.float32)/255
48
+ for j in (z-1,z,z+1)])
49
+ tensor=torch.from_numpy(arr.copy())[None]
50
+ corrupted,target,mask=module.corrupt_center(tensor)
51
+ m=mask.numpy()[0,0]
52
+ original=target.numpy()[0,0]
53
+ changed=(corrupted[:,1:2].numpy()[0,0]!=original)
54
+ selected=int(m.sum())
55
+ changed_count=int((changed&m).sum())
56
+ rows.append(dict(sample=i,z=z,masked=selected,changed=changed_count,
57
+ masked_unchanged=selected-changed_count,
58
+ patch_std=float(original.std()),
59
+ patch_mean=float(original.mean()),
60
+ masked_target_var=float(original[m].var()) if selected else 0.0))
61
+ with (out/'patch_audit.csv').open('w',newline='') as f:
62
+ w=csv.DictWriter(f,fieldnames=rows[0].keys());w.writeheader();w.writerows(rows)
63
+ total_mask=sum(r['masked'] for r in rows)
64
+ unchanged=sum(r['masked_unchanged'] for r in rows)
65
+ print(f'Masked target pixels unchanged after replacement: {unchanged}/{total_mask} ({unchanged/max(1,total_mask):.2%})')
66
+ print('Patch std: min/median/max:',*[round(v,5) for v in
67
+ (min(r['patch_std'] for r in rows),np.median([r['patch_std'] for r in rows]),max(r['patch_std'] for r in rows))])
68
+ # Counterfactual probe: keep corrupted center fixed, change only neighboring channels.
69
+ # Measures neighbor dependence, NOT proof of leakage or restoration quality.
70
+ z=5;h,w=ims[z].shape;y=(h-a.patch)//2;x=(w-a.patch)//2
71
+ arr=np.stack([np.asarray(ims[j][y:y+a.patch,x:x+a.patch],dtype=np.float32)/255
72
+ for j in (z-1,z,z+1)])
73
+ base=torch.from_numpy(arr.copy())[None]
74
+ corrupted,target,mask=module.corrupt_center(base)
75
+ with torch.inference_mode():
76
+ baseline=model(corrupted)
77
+ zero=corrupted.clone();zero[:,0]=0;zero[:,2]=0
78
+ same=corrupted.clone();same[:,0]=same[:,1];same[:,2]=same[:,1]
79
+ swap=corrupted.clone();swap[:,0],swap[:,2]=corrupted[:,2].clone(),corrupted[:,0].clone()
80
+ delta_zero=(model(zero)-baseline).abs().mean().item()
81
+ delta_same=(model(same)-baseline).abs().mean().item()
82
+ delta_swap=(model(swap)-baseline).abs().mean().item()
83
+ print('Neighbor counterfactual mean |output change|:')
84
+ print(' zero neighbors:',round(delta_zero,6))
85
+ print(' copy corrupted center into neighbors:',round(delta_same,6))
86
+ print(' swap neighbors:',round(delta_swap,6))
87
+ report={
88
+ 'checkpoint':str(a.checkpoint),'mode':ck['mode'],
89
+ 'train_files':[f.name for f in train],'heldout_files':[f.name for f in test],
90
+ 'filename_overlap':bool(set(train)&set(test)),
91
+ 'masked_pixels':total_mask,'masked_pixels_unchanged':unchanged,
92
+ 'masked_unchanged_fraction':unchanged/max(1,total_mask),
93
+ 'patch_std_min':min(r['patch_std'] for r in rows),
94
+ 'patch_std_median':float(np.median([r['patch_std'] for r in rows])),
95
+ 'neighbor_probe_mean_abs_change':{'zero':delta_zero,'copy_center':delta_same,'swap':delta_swap},
96
+ 'limitations':[
97
+ 'This is a code and behavior audit, not proof of no information leakage.',
98
+ 'Adjacent slice channels remain unmasked; correlated neighboring pixels can reveal central structure.',
99
+ 'Random neighboring-pixel replacement can equal the target by chance.',
100
+ 'Training/test adjacent sections are not independent specimens.',
101
+ 'No retraining or performance assessment is performed.'
102
+ ]}
103
+ (out/'audit.json').write_text(json.dumps(report,indent=2))
104
+ print('Saved:',out/'audit.json',out/'patch_audit.csv')
105
+ if __name__=='__main__': main()
@@ -0,0 +1,64 @@
1
+ """Command-line dispatcher for the existing StructRestore-3D research scripts.
2
+
3
+ All data and checkpoint paths are relative to the shell working directory.
4
+ """
5
+ from pathlib import Path
6
+ import runpy
7
+ import sys
8
+
9
+ from . import __version__
10
+
11
+ # Command -> packaged script name, positional arguments prepended for existing CLI.
12
+ COMMANDS = {
13
+ "audit-v1": ("train_v1_audited", ["audit"]),
14
+ "train-v1": ("train_v1_audited", ["train"]),
15
+ "compare": ("compare_v0_v1", []),
16
+ "diagnose": ("diagnose_v1_inference", []),
17
+ "mask-audit": ("audit_masking", []),
18
+ "train-single-v0": ("train_fullres", []),
19
+ "train-25d-v0": ("structrestore_25d", []),
20
+ "train-baseline": ("train_baseline", []),
21
+ "infer-baseline": ("infer_baseline", []),
22
+ "validate-25d": ("validate_25d", []),
23
+ "register-probe": ("step29_registration_probe", []),
24
+ }
25
+
26
+
27
+ def main():
28
+ if len(sys.argv) < 2 or sys.argv[1] in ("-h", "--help", "help"):
29
+ print(f"StructRestore-3D {__version__} (experimental research software)")
30
+ print("Usage: structrestore3d COMMAND [OPTIONS]")
31
+ print("Commands:")
32
+ for command in COMMANDS:
33
+ print(f" {command}")
34
+ print("\nExample: structrestore3d audit-v1")
35
+ print("Run 'structrestore3d train-v1 --help' for training arguments.")
36
+ print("All image data and trained model checkpoints must be supplied separately.")
37
+ return 0
38
+ if sys.argv[1] in ("-V", "--version", "version"):
39
+ print(__version__)
40
+ return 0
41
+ command = sys.argv[1]
42
+ if command not in COMMANDS:
43
+ print(f"Unknown command: {command!r}; run structrestore3d --help", file=sys.stderr)
44
+ return 2
45
+ script, prefix = COMMANDS[command]
46
+ script_path = Path(__file__).resolve().parent / f"{script}.py"
47
+ if not script_path.is_file():
48
+ print(f"The packaged script {script} is missing. Run python tools/sync_package.py before building.", file=sys.stderr)
49
+ return 1
50
+ # Original research scripts import their siblings as top-level modules.
51
+ # Keep backward compatibility without altering the user's scripts.
52
+ package_dir = str(script_path.parent)
53
+ if package_dir not in sys.path:
54
+ sys.path.insert(0, package_dir)
55
+ arguments = sys.argv[2:]
56
+ if command == "mask-audit" and "--script" not in arguments:
57
+ arguments += ["--script", str(script_path.parent / "structrestore_25d.py")]
58
+ sys.argv = [str(script_path), *prefix, *arguments]
59
+ runpy.run_path(str(script_path), run_name="__main__")
60
+ return 0
61
+
62
+
63
+ if __name__ == "__main__":
64
+ raise SystemExit(main())
@@ -0,0 +1,213 @@
1
+ #!/usr/bin/env python3
2
+ """Step 27: controlled V0 vs V1 EM FIB-SEM added-noise stress test.
3
+
4
+ Reference: original *naturally noisy* TIFF, NOT verified clean ground truth.
5
+ V1: masked-partition prediction: each pixel is inferred ONLY while concealed
6
+ at the same pixel coordinate in all input sections, avoiding unmasked
7
+ inference mismatch. Masked pixels are given independent uniform filler.
8
+ V0: native inference exactly as originally used. Comparisons confound training
9
+ changes and inference protocols; not causal architecture ablations.
10
+ No files in existing results directories are written or altered.
11
+ """
12
+ import argparse, csv, json, sys, time
13
+ from pathlib import Path
14
+ import numpy as np
15
+ import tifffile
16
+ import torch
17
+ from scipy.ndimage import gaussian_filter, median_filter
18
+ from skimage.metrics import structural_similarity, peak_signal_noise_ratio
19
+ import matplotlib
20
+ matplotlib.use('Agg')
21
+ import matplotlib.pyplot as plt
22
+
23
+ # Loads existing project architectures without modifying source or checkpoints.
24
+ SCRIPT_DIR=Path(__file__).resolve().parent
25
+ if SCRIPT_DIR.name!='scripts' and (SCRIPT_DIR/'scripts').is_dir():
26
+ SCRIPT_DIR=SCRIPT_DIR/'scripts'
27
+ sys.path.insert(0,str(SCRIPT_DIR))
28
+ from train_fullres import TinyCNN
29
+ from structrestore_25d import Net25D
30
+ from train_v1_audited import Model as V1Model
31
+
32
+ DEFAULT_PATHS={
33
+ 'v0_single':'results/baseline_fullres/baseline.pt',
34
+ 'v0_plain':'results/25d_plain/model.pt',
35
+ 'v0_gated':'results/25d_gated/model.pt',
36
+ 'v1_single':'results/v1_single/model.pt',
37
+ 'v1_plain':'results/v1_plain/model.pt',
38
+ 'v1_gated':'results/v1_gated/model.pt',
39
+ }
40
+
41
+ def load_models(paths,dev):
42
+ models={}
43
+ for key,path in paths.items():
44
+ ck=torch.load(path,map_location='cpu',weights_only=True)
45
+ mode=key.split('_')[1]
46
+ if key.startswith('v0'):
47
+ net=TinyCNN(width=int(ck.get('width',32))) if mode=='single' else Net25D(mode,int(ck.get('width',32)))
48
+ if mode!='single' and ck.get('mode')!=mode:raise ValueError(f'V0 mode mismatch: {key}: {ck.get("mode")}')
49
+ else:
50
+ if ck.get('mode')!=mode:raise ValueError(f'V1 mode mismatch: {key}: {ck.get("mode")}')
51
+ net=V1Model(mode,int(ck.get('width',32)))
52
+ net.load_state_dict(ck['state_dict'],strict=True)
53
+ models[key]=net.to(dev).eval()
54
+ print('Loaded',key, '<-',path,flush=True)
55
+ return models
56
+
57
+ def v0_predict(model,mode,stack,dev,pad):
58
+ arr=stack[1:2] if mode=='single' else stack
59
+ with torch.inference_mode():
60
+ inp=torch.from_numpy(np.ascontiguousarray(arr[None])).to(dev)
61
+ out=model(inp).clamp(0,1)[0,0].cpu().numpy()
62
+ return out[pad:-pad,pad:-pad]
63
+
64
+ def masked_partition_v1(model,stack,dev,pad,groups,rng):
65
+ """Predict every pixel *once* under a disjoint random ~5% mask partition.
66
+
67
+ Independent uniform replacement for each masked channel matches training
68
+ mask value distribution. All channels use exactly the same boolean mask.
69
+ """
70
+ _,h,w=stack.shape
71
+ group_ids=rng.integers(0,groups,size=(h,w),dtype=np.uint8 if groups<=256 else np.int16)
72
+ output=np.empty((h-2*pad,w-2*pad),dtype=np.float32)
73
+ counts=np.zeros(output.shape,dtype=np.uint8)
74
+ base=torch.from_numpy(np.ascontiguousarray(stack[None])).to(dev)
75
+ with torch.inference_mode():
76
+ for g in range(groups):
77
+ selected=(group_ids==g)
78
+ if not np.any(selected[pad:-pad,pad:-pad]):continue
79
+ mask=torch.from_numpy(selected[None,None]).to(dev)
80
+ # Draw filler per channel independently, as in V1 training.
81
+ filler=torch.rand_like(base)
82
+ masked=torch.where(mask.expand_as(base),filler,base)
83
+ prediction=model(masked).clamp(0,1)[0,0].cpu().numpy()[pad:-pad,pad:-pad]
84
+ interior=selected[pad:-pad,pad:-pad]
85
+ output[interior]=prediction[interior]
86
+ counts[interior]+=1
87
+ if not np.all(counts==1):raise RuntimeError('Mask partitions did not cover each interior pixel exactly once')
88
+ return output
89
+
90
+ def calc_metrics(ref,img):
91
+ img=np.clip(img,0,1).astype(np.float32)
92
+ gy,gx=np.gradient(ref)
93
+ py,px=np.gradient(img)
94
+ return {'psnr':float(peak_signal_noise_ratio(ref,img,data_range=1.0)),
95
+ 'ssim':float(structural_similarity(ref,img,data_range=1.0)),
96
+ 'edge_rmse':float(np.sqrt(np.mean((gy-py)**2+(gx-px)**2)))}
97
+
98
+ def main():
99
+ p=argparse.ArgumentParser(description=__doc__)
100
+ p.add_argument('--input',default='data/raw')
101
+ p.add_argument('--out',default='results/v0_v1_comparison')
102
+ p.add_argument('--crop',type=int,default=256,help='Central output crop width; extra 16px context used')
103
+ p.add_argument('--trials',type=int,default=2)
104
+ p.add_argument('--sigmas',type=float,nargs='+',default=[0.05,0.10])
105
+ p.add_argument('--groups',type=int,default=20,help='V1 disjoint random mask groups, 20 ~ 5%% masked per pass')
106
+ p.add_argument('--seed',type=int,default=2026)
107
+ p.add_argument('--models-root',default='.',help='Project root containing model checkpoints')
108
+ args=p.parse_args()
109
+ pad=16
110
+ if args.crop<64 or args.trials<1 or args.groups<2 or args.groups>255 or any(not 0<s<0.5 for s in args.sigmas):
111
+ p.error('Require crop>=64, trials>=1, 2<=groups<=255 and 0<sigma<0.5')
112
+ files=sorted(Path(args.input).glob('slice_*.tif'))
113
+ if len(files)!=16:raise ValueError(f'Expected 16 original TIFFs; found {len(files)}')
114
+ paths={k:Path(args.models_root)/v for k,v in DEFAULT_PATHS.items()}
115
+ missing=[str(f) for f in paths.values() if not f.is_file()]
116
+ if missing:raise FileNotFoundError('Missing checkpoints: '+', '.join(missing))
117
+ out=Path(args.out)
118
+ if out.exists() and any(out.iterdir()):
119
+ raise FileExistsError(f'Output directory is not empty: {out}. Use new --out to avoid overwriting experiments.')
120
+ out.mkdir(parents=True,exist_ok=True)
121
+ dev=torch.device('mps' if torch.backends.mps.is_available() else 'cpu')
122
+ print('Device:',dev,flush=True)
123
+ models=load_models(paths,dev)
124
+ rng=np.random.default_rng(args.seed)
125
+ rows=[];figure_done=False
126
+ # Require both context sections to belong to the four held-out slices.
127
+ for z in (13,14):
128
+ ims=[tifffile.memmap(str(files[j])) for j in (z-1,z,z+1)]
129
+ h,w=ims[1].shape;side=args.crop
130
+ if min(h,w)<side+2*pad:raise ValueError('Image too small for crop and 16px halo')
131
+ centers=[(h//4,w//4),(h//2,w//2),(3*h//4,3*w//4)]
132
+ for region,(cy,cx) in enumerate(centers):
133
+ y=max(pad,min(h-side-pad,cy-side//2))
134
+ x=max(pad,min(w-side-pad,cx-side//2))
135
+ stack=np.stack([np.asarray(im[y-pad:y+side+pad,x-pad:x+side+pad],dtype=np.float32)/255 for im in ims])
136
+ reference=stack[1,pad:-pad,pad:-pad].copy()
137
+ for sigma in args.sigmas:
138
+ for trial in range(args.trials):
139
+ noisy=np.clip(stack+rng.normal(0,sigma,stack.shape).astype(np.float32),0,1)
140
+ noisy_center=noisy[1,pad:-pad,pad:-pad]
141
+ candidates={'added_noise':noisy_center,
142
+ 'gaussian_0p8':gaussian_filter(noisy_center,0.8),
143
+ 'median_3x3':median_filter(noisy_center,size=3)}
144
+ times={k:0.0 for k in candidates}
145
+ # SHARE one mask assignment and filler across the three V1 networks:
146
+ # reset torch generator before each call so masks/filler match.
147
+ # This probe is repeatable on CPU and MPS within each run.
148
+ rng_state=rng.bit_generator.state
149
+ for key,model in models.items():
150
+ start=time.perf_counter()
151
+ if key.startswith('v0'):
152
+ candidates[key]=v0_predict(model,key.split('_')[1],noisy,dev,pad)
153
+ else:
154
+ rng.bit_generator.state=rng_state
155
+ torch.manual_seed(args.seed+100000*z+1000*region+100*trial+int(round(100*sigma)))
156
+ candidates[key]=masked_partition_v1(model,noisy,dev,pad,args.groups,rng)
157
+ times[key]=time.perf_counter()-start
158
+ # Advance RNG state by exactly one V1 mask draw for next noise trial.
159
+ for name,img in candidates.items():
160
+ m=calc_metrics(reference,img)
161
+ rows.append({'slice':f'slice_{z:04d}','region':region,'sigma':sigma,
162
+ 'trial':trial,'method':name,'seconds':round(times[name],5),**m})
163
+ if not figure_done:
164
+ fig,axes=plt.subplots(3,3,figsize=(12,12),constrained_layout=True)
165
+ ordered=[('Reference (native noise)',reference),('Added noise',noisy_center),
166
+ ('Gaussian 0.8',candidates['gaussian_0p8']),
167
+ ('V0 single',candidates['v0_single']),('V0 plain',candidates['v0_plain']),
168
+ ('V0 gated',candidates['v0_gated']),('V1 single',candidates['v1_single']),
169
+ ('V1 plain',candidates['v1_plain']),('V1 gated',candidates['v1_gated'])]
170
+ for ax,(name,img) in zip(axes.ravel(),ordered):
171
+ ax.imshow(img,cmap='gray',vmin=0,vmax=1,interpolation='nearest')
172
+ ax.set_title(name);ax.axis('off')
173
+ fig.suptitle(f'Added noise sigma {sigma} | {files[z].stem} | region {region}; reference is NOT clean')
174
+ fig.savefig(out/'comparison.png',dpi=140)
175
+ plt.close(fig);figure_done=True
176
+ print('Finished',files[z].stem,flush=True)
177
+ methods=['added_noise','gaussian_0p8','median_3x3']+list(DEFAULT_PATHS)
178
+ with (out/'per_trial.csv').open('w',newline='') as f:
179
+ writer=csv.DictWriter(f,fieldnames=list(rows[0]));writer.writeheader();writer.writerows(rows)
180
+ summary=[]
181
+ for sigma in args.sigmas:
182
+ for method in methods:
183
+ subset=[r for r in rows if r['sigma']==sigma and r['method']==method]
184
+ r={'sigma':sigma,'method':method,'n_correlated_evaluations':len(subset)}
185
+ for metric in ['psnr','ssim','edge_rmse','seconds']:
186
+ vals=np.array([v[metric] for v in subset],dtype=np.float64)
187
+ r[metric+'_mean']=round(float(vals.mean()),6)
188
+ r[metric+'_sd']=round(float(vals.std(ddof=1)),6) if len(vals)>1 else 0.
189
+ summary.append(r)
190
+ with (out/'summary.csv').open('w',newline='') as f:
191
+ writer=csv.DictWriter(f,fieldnames=list(summary[0]));writer.writeheader();writer.writerows(summary)
192
+ metadata={'experiment':'V0 vs V1 / added-noise stress test','device':str(dev),'seed':args.seed,
193
+ 'train_test_boundary':'train slices 0000-0011, heldout 0012-0015',
194
+ 'evaluated_slice_centers':['slice_0013','slice_0014'],
195
+ 'slice_context_always_heldout':True,'crop':args.crop,'trials':args.trials,
196
+ 'sigmas':args.sigmas,'v1_mask_groups':args.groups,
197
+ 'checkpoints':{k:str(v) for k,v in paths.items()},
198
+ 'limitations':[
199
+ 'Original FIB-SEM reference retains unknown native noise; it is NOT clean ground truth.',
200
+ 'Metrics assess fidelity to original after extra simulated noise, not removal of native noise.',
201
+ 'V1 inference predicts every pixel once while directly masked using random disjoint partitions; V0 uses native unmasked inference.',
202
+ 'V1 changes training masking and patch sampling plus inference procedure: not a causal isolation of any one change.',
203
+ 'No independent biological specimens; slices, patches and noise replicates are correlated.',
204
+ 'Edge RMSE is gradient fidelity to noisy original, not verified membrane-preservation accuracy.',
205
+ 'Models were trained on native noisy images; extra Gaussian noise is an out-of-distribution stress test.'
206
+ ]}
207
+ (out/'metadata.json').write_text(json.dumps(metadata,indent=2))
208
+ print('\nMEAN METRICS: high PSNR/SSIM is better; low edge_RMSE is better',flush=True)
209
+ for r in summary:
210
+ print(f"sigma={r['sigma']:.2f} {r['method']:<14} PSNR={r['psnr_mean']:6.2f} SSIM={r['ssim_mean']:.4f} edge_RMSE={r['edge_rmse_mean']:.4f}",flush=True)
211
+ print('\nSaved:',out/'summary.csv',out/'per_trial.csv',out/'comparison.png',out/'metadata.json',flush=True)
212
+
213
+ if __name__=='__main__':main()