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.
- structrestore_3d-0.1.0a1/LICENSE +21 -0
- structrestore_3d-0.1.0a1/PKG-INFO +80 -0
- structrestore_3d-0.1.0a1/README.md +55 -0
- structrestore_3d-0.1.0a1/pyproject.toml +39 -0
- structrestore_3d-0.1.0a1/setup.cfg +4 -0
- structrestore_3d-0.1.0a1/src/structrestore3d/__init__.py +3 -0
- structrestore_3d-0.1.0a1/src/structrestore3d/__main__.py +4 -0
- structrestore_3d-0.1.0a1/src/structrestore3d/audit_masking.py +105 -0
- structrestore_3d-0.1.0a1/src/structrestore3d/cli.py +64 -0
- structrestore_3d-0.1.0a1/src/structrestore3d/compare_v0_v1.py +213 -0
- structrestore_3d-0.1.0a1/src/structrestore3d/diagnose_v1_inference.py +260 -0
- structrestore_3d-0.1.0a1/src/structrestore3d/infer_baseline.py +88 -0
- structrestore_3d-0.1.0a1/src/structrestore3d/step29_registration_probe.py +244 -0
- structrestore_3d-0.1.0a1/src/structrestore3d/structrestore_25d.py +150 -0
- structrestore_3d-0.1.0a1/src/structrestore3d/train_baseline.py +75 -0
- structrestore_3d-0.1.0a1/src/structrestore3d/train_fullres.py +78 -0
- structrestore_3d-0.1.0a1/src/structrestore3d/train_v1_audited.py +130 -0
- structrestore_3d-0.1.0a1/src/structrestore3d/validate_25d.py +120 -0
- structrestore_3d-0.1.0a1/src/structrestore_3d.egg-info/PKG-INFO +80 -0
- structrestore_3d-0.1.0a1/src/structrestore_3d.egg-info/SOURCES.txt +22 -0
- structrestore_3d-0.1.0a1/src/structrestore_3d.egg-info/dependency_links.txt +1 -0
- structrestore_3d-0.1.0a1/src/structrestore_3d.egg-info/entry_points.txt +2 -0
- structrestore_3d-0.1.0a1/src/structrestore_3d.egg-info/requires.txt +6 -0
- 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,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()
|