0)
+ #stel=np.array([[1, 1],[1,1]]).astype(bool)
+ stel=np.array([[1,1,1], [1,1,1], [1,1,1]]).astype(bool)
+ #stel=np.array([[1,1,1,1],[1, 1, 1, 1],[1,1,1,1],[1,1,1,1]]).astype(bool)
+ gg3gd=np.zeros(gg3g.shape)
+ for i in range(gg3g.shape[0]):
+ gg3gd[i,:,:]=binary_dilation(gg3g[i,:,:],structure=stel,iterations=iter_num)
+ out = gg3.copy()
+ out[gg3gd==1]=0
+ return out
+
+def markInvalid(seg, iter_num=2, do_2d=True):
+ # find invalid
+ # if do erosion(seg==0), then miss the border
+ if do_2d:
+ stel=np.array([[1,1,1], [1,1,1]]).astype(bool)
+ if len(seg.shape)==2:
+ out = binary_dilation(seg>0, structure=stel, iterations=iter_num)
+ seg[out==0] = -1
+ else: # save memory
+ for z in range(seg.shape[0]):
+ tmp = seg[z] # by reference
+ out = binary_dilation(tmp>0, structure=stel, iterations=iter_num)
+ tmp[out==0] = -1
+ else:
+ stel=np.array([[1,1,1], [1,1,1], [1,1,1]]).astype(bool)
+ out = binary_dilation(seg>0, structure=stel, iterations=iter_num)
+ seg[out==0] = -1
+ return seg
diff --git a/legacy/Pretrain/utils/show.py b/legacy/Pretrain/utils/show.py
new file mode 100644
index 0000000..17ce46b
--- /dev/null
+++ b/legacy/Pretrain/utils/show.py
@@ -0,0 +1,299 @@
+import os
+import math
+import numpy as np
+from PIL import Image
+
+def show(img3d):
+ # only used for image with shape [18, 160, 160, 3]
+ num = img3d.shape[0]
+ column = 5
+ row = math.ceil(num / float(column))
+ size = img3d.shape[1]
+ img_all = np.zeros((size*row, size*column, 3), dtype=np.uint8)
+ for i in range(row):
+ for j in range(column):
+ index = i*column + j
+ if index >= num:
+ img = np.zeros_like(img3d[0], dtype=np.uint8)
+ else:
+ if np.max(img3d[index]) > 1:
+ img = (img3d[index]).astype(np.uint8)
+ else:
+ img = (img3d[index] * 255).astype(np.uint8)
+ img_all[i*size:(i+1)*size, j*size:(j+1)*size] = img
+ return img_all
+
+def show_one(img3d):
+ # only used for image with shape [18, 160, 160]
+ num = img3d.shape[0]
+ column = 5
+ row = math.ceil(num / float(column))
+ size = img3d.shape[1]
+ img_all = np.zeros((size*row, size*column), dtype=np.uint8)
+ for i in range(row):
+ for j in range(column):
+ index = i*column + j
+ if index >= num:
+ img = np.zeros_like(img3d[0], dtype=np.uint8)
+ else:
+ img = (img3d[index] * 255).astype(np.uint8)
+ img_all[i*size:(i+1)*size, j*size:(j+1)*size] = img
+ return img_all
+
+
+def show_one_(img3d):
+ num = img3d.shape[0]
+ column = 5
+ row = math.ceil(num / float(column))
+ size = img3d.shape[1]
+ img_all = np.zeros((size*row, size*column), dtype=np.uint8)
+ for i in range(row):
+ for j in range(column):
+ index = i*column + j
+ if index >= num:
+ img = np.zeros_like(img3d[0], dtype=np.uint8)
+ else:
+ img = (img3d[index]).astype(np.uint8)
+ img_all[i*size:(i+1)*size, j*size:(j+1)*size] = img
+ return img_all
+
+def show_CE(img3d):
+ # only used for image with shape [18, 160, 160]
+ num = img3d.shape[0]
+ column = 5
+ row = math.ceil(num / float(column))
+ size = img3d.shape[1]
+ img_all = np.zeros((size*row, size*column), dtype=np.uint8)
+ for i in range(row):
+ for j in range(column):
+ index = i*column + j
+ if index >= num:
+ img = np.zeros_like(img3d[0], dtype=np.uint8)
+ else:
+ img = (img3d[index]).astype(np.uint8)
+ img_all[i*size:(i+1)*size, j*size:(j+1)*size] = img
+ return img_all
+
+def training_show(iters, inputs, label, pred_bound, cache_path, if_skele=None, skele=None, pred_skele=None):
+ img_input = np.repeat(inputs[0].data.cpu().numpy(), 3, 0)
+ img_input = np.transpose(img_input, (1,2,3,0))
+ img_input = show(img_input)
+ input_placehplder = np.zeros_like(img_input, dtype=np.uint8)
+ im_cat1 = np.concatenate([img_input, input_placehplder], axis=1)
+
+ img_label = label[0][0:3].data.cpu().numpy()
+ img_label = np.transpose(img_label, (1,2,3,0))
+ img_label = show(img_label)
+
+ img_pred_bound = pred_bound[0][0:3].data.cpu().numpy()
+ img_pred_bound = np.transpose(img_pred_bound, (1,2,3,0))
+ img_pred_bound = show(img_pred_bound)
+ im_cat2 = np.concatenate([img_pred_bound, img_label], axis=1)
+
+ if if_skele is not None:
+ img_skele = np.repeat(skele[0, 0:1].data.cpu().numpy(), 3, 0)
+ img_skele = np.transpose(img_skele, (1,2,3,0))
+ img_skele = show(img_skele)
+
+ img_pred_skele = np.repeat(pred_skele[0, 0:1].data.cpu().numpy(), 3, 0)
+ img_pred_skele = np.transpose(img_pred_skele, (1,2,3,0))
+ img_pred_skele = show(img_pred_skele)
+ im_cat3 = np.concatenate([img_pred_skele, img_skele], axis=1)
+
+ im_cat = np.concatenate([im_cat1, im_cat2, im_cat3], axis=0)
+ else:
+ im_cat = np.concatenate([im_cat1, im_cat2], axis=0)
+ Image.fromarray(im_cat).save(os.path.join(cache_path, '%06d.png' % iters))
+
+def training_show_pretrain(iters, pred, label, cache_path, loss_mode='CrossEntropy'):
+ img_input = pred[0].data.cpu().numpy()
+ if loss_mode == 'CrossEntropy':
+ img_input = show_CE(img_input)
+ else:
+ img_input[img_input < 0] = 0
+ img_input[img_input > 1] = 1
+ img_input = show_one(img_input)
+ img_label = label[0].data.cpu().numpy()
+ img_label = show_one(img_label)
+ im_cat = np.concatenate([img_input, img_label], axis=1)
+ Image.fromarray(im_cat).save(os.path.join(cache_path, '%06d.png' % iters))
+
+
+def show_inpaining(iters, pred, label, mask, cache_path):
+ pred = pred[0].data.cpu().numpy()
+ label = label[0].data.cpu().numpy()
+ mask = mask[0].data.cpu().numpy()
+ inputs = label * mask
+ inputs = np.squeeze(inputs)
+ pred = np.squeeze(pred)
+ inputs = inputs[14:-14, 106:-106, 106:-106]
+ pred[pred < 0] = 0; pred[pred > 1] =1
+ pred_img = show_one(pred)
+ inputs_img = show_one(inputs)
+ im_cat = np.concatenate([inputs_img, pred_img], axis=1)
+ Image.fromarray(im_cat).save(os.path.join(cache_path, '%06d.png' % iters))
+
+
+def show_affs(iters, inputs, pred, target, cache_path, model_type='mala'):
+ pred = pred[0].data.cpu().numpy()
+ inputs = inputs[0].data.cpu().numpy()
+ target = target[0].data.cpu().numpy()
+ inputs = np.squeeze(inputs)
+ if model_type == 'mala':
+ inputs = inputs[14:-14, 106:-106, 106:-106]
+ inputs = inputs[:,:,:,np.newaxis]
+ inputs = np.repeat(inputs, 3, 3)
+ pred = np.transpose(pred, (1,2,3,0))
+ target = np.transpose(target, (1,2,3,0))
+ inputs[inputs<0]=0; inputs[inputs>1]=1
+ pred[pred<0]=0; pred[pred>1]=1
+ target[target<0]=0; target[target>1]=1
+ inputs_img = show(inputs)
+ pred_img = show(pred)
+ target_img = show(target)
+ im_cat = np.concatenate([inputs_img, pred_img, target_img], axis=1)
+ Image.fromarray(im_cat).save(os.path.join(cache_path, '%06d.png' % iters))
+
+def show_bound(iters, inputs, pred, target, cache_path, model_type='mala'):
+ pred = pred[0].data.cpu().numpy()
+ inputs = inputs[0].data.cpu().numpy()
+ target = target[0].data.cpu().numpy()
+
+ inputs = np.squeeze(inputs)
+ inputs = inputs[:,:,:,np.newaxis]
+ inputs = np.repeat(inputs, 3, 3)
+
+ pred = np.squeeze(pred)
+ pred = pred[:,:,:,np.newaxis]
+ pred = np.repeat(pred, 3, 3)
+
+ target = np.squeeze(target)
+ target = target[:,:,:,np.newaxis]
+ target = np.repeat(target, 3, 3)
+
+ inputs[inputs<0]=0; inputs[inputs>1]=1
+ pred[pred<0]=0; pred[pred>1]=1
+ target[target<0]=0; target[target>1]=1
+ inputs_img = show(inputs)
+ pred_img = show(pred)
+ target_img = show(target)
+ im_cat = np.concatenate([inputs_img, pred_img, target_img], axis=1)
+ Image.fromarray(im_cat).save(os.path.join(cache_path, '%06d.png' % iters))
+
+def show_bound2(iters, inputs, pred, target, cache_path, model_type='mala'):
+ pred = pred[0].data.cpu().numpy()
+ inputs = inputs[0].data.cpu().numpy()
+ target = target[0].data.cpu().numpy()
+
+ inputs = np.squeeze(inputs)
+ inputs = inputs[:,:,:,np.newaxis]
+ inputs = np.repeat(inputs, 3, 3)
+
+ pred = np.squeeze(pred)
+ pred = pred[:,:,:,np.newaxis]
+ pred = np.repeat(pred, 3, 3)
+
+ target = np.squeeze(target)
+ target = target[:,:,:,np.newaxis]
+ target = np.repeat(target, 3, 3)
+
+ inputs[inputs<0]=0; inputs[inputs>1]=1
+ pred[pred<0]=0; pred[pred>1]=1
+ target[target<0]=0; target[target>1]=1
+ inputs_img = show(inputs)
+ pred_img = show(pred)
+ target_img = show(target)
+ im_cat = np.concatenate([inputs_img, pred_img, target_img], axis=1)
+ Image.fromarray(im_cat).save(os.path.join(cache_path, '%06d_hog.png' % iters))
+
+def class_color(lb):
+ d, h, w = lb.shape
+ lb_color1 = np.zeros((d, h, w), dtype=np.uint8)
+ lb_color2 = np.zeros((d, h, w), dtype=np.uint8)
+ lb_color3 = np.zeros((d, h, w), dtype=np.uint8)
+ ids_0 = lb == 0
+ ids_1 = lb == 1
+ lb_color1[ids_0] = 0; lb_color2[ids_0] = 0; lb_color3[ids_0] = 255
+ lb_color1[ids_1] = 0; lb_color2[ids_1] = 255; lb_color3[ids_1] = 0
+ lb_color = np.concatenate([lb_color1[:,:,:,np.newaxis], lb_color2[:,:,:,np.newaxis], lb_color3[:,:,:,np.newaxis]], axis=3)
+ return lb_color
+
+
+def show_affs_pseudo(iters, inputs, pred, target, mask, cache_path, model_type='mala'):
+ pred = pred[0].data.cpu().numpy()
+ inputs = inputs[0].data.cpu().numpy()
+ target = target[0].data.cpu().numpy()
+ mask = mask[0].data.cpu().numpy()
+ inputs = np.squeeze(inputs)
+ if model_type == 'mala':
+ inputs = inputs[14:-14, 106:-106, 106:-106]
+ inputs = inputs[:,:,:,np.newaxis]
+ inputs = np.repeat(inputs, 3, 3)
+ pred = np.transpose(pred, (1,2,3,0))
+ target = np.transpose(target, (1,2,3,0))
+ affs_z = class_color(target[:, :, :, 0]) * mask[0][:,:,:,np.newaxis]
+ affs_y = class_color(target[:, :, :, 1]) * mask[1][:,:,:,np.newaxis]
+ affs_x = class_color(target[:, :, :, 2]) * mask[2][:,:,:,np.newaxis]
+ inputs_img = show(inputs)
+ pred_img = show(pred)
+ # target_img = show(target)
+ mask = np.transpose(mask, (1,2,3,0))
+ mask_img = show(mask)
+ affs_z_img = show(affs_z)
+ affs_y_img = show(affs_y)
+ affs_x_img = show(affs_x)
+ # im_cat = np.concatenate([inputs_img, pred_img, target_img], axis=1)
+ im_cat1 = np.concatenate([inputs_img, pred_img], axis=1)
+ im_cat2 = np.concatenate([mask_img, affs_z_img], axis=1)
+ im_cat3 = np.concatenate([affs_y_img, affs_x_img], axis=1)
+ im_cat = np.concatenate([im_cat1, im_cat2, im_cat3], axis=0)
+ Image.fromarray(im_cat).save(os.path.join(cache_path, '%06d.png' % iters))
+
+def show_affs_whole(iters, out_affs, gt_affs, cache_path, index):
+ out_affs = out_affs[:, -1, ...]
+ gt_affs = gt_affs[:, -1, ...]
+ out_affs = (out_affs * 255).astype(np.uint8)
+ out_affs = np.transpose(out_affs, (1,2,0))
+ gt_affs = (gt_affs * 255).astype(np.uint8)
+ gt_affs = np.transpose(gt_affs, (1,2,0))
+ im_cat = np.concatenate([out_affs, gt_affs], axis=1)
+ Image.fromarray(im_cat).save(os.path.join(cache_path, '%06d_%d.png' % (iters,index)))
+
+def show_bound_whole(iters, out_affs, gt_affs, cache_path, index):
+ out_affs = out_affs.squeeze()[0]
+ out_affs = (out_affs * 255).astype(np.uint8)
+ gt_affs = gt_affs.squeeze()[0]
+ gt_affs = (gt_affs * 255).astype(np.uint8)
+ im_cat = np.concatenate([out_affs, gt_affs], axis=1)
+ Image.fromarray(im_cat).save(os.path.join(cache_path, '%06d_%d.png' % (iters,index)))
+
+def show_affs_consistency(iters, inputs, pred, target, inputs_u, out_u1, out_u2, cache_path):
+ pred = pred[0].data.cpu().numpy()
+ inputs = inputs[0].data.cpu().numpy()
+ target = target[0].data.cpu().numpy()
+ inputs = np.squeeze(inputs)
+ inputs = inputs[14:-14, 106:-106, 106:-106]
+ inputs = inputs[:,:,:,np.newaxis]
+ inputs = np.repeat(inputs, 3, 3)
+ pred = np.transpose(pred, (1,2,3,0))
+ target = np.transpose(target, (1,2,3,0))
+ inputs_img = show(inputs)
+ pred_img = show(pred)
+ target_img = show(target)
+ im_cat1 = np.concatenate([inputs_img, pred_img, target_img], axis=1)
+
+ out_u1 = out_u1[0].data.cpu().numpy()
+ inputs_u = inputs_u[0].data.cpu().numpy()
+ out_u2 = out_u2[0].data.cpu().numpy()
+ inputs_u = np.squeeze(inputs_u)
+ inputs_u = inputs_u[14:-14, 106:-106, 106:-106]
+ inputs_u = inputs_u[:,:,:,np.newaxis]
+ inputs_u = np.repeat(inputs_u, 3, 3)
+ out_u1 = np.transpose(out_u1, (1,2,3,0))
+ out_u2 = np.transpose(out_u2, (1,2,3,0))
+ inputs_u_img = show(inputs_u)
+ out_u1_img = show(out_u1)
+ out_u2_img = show(out_u2)
+ im_cat2 = np.concatenate([inputs_u_img, out_u1_img, out_u2_img], axis=1)
+ im_cat = np.concatenate([im_cat1, im_cat2], axis=0)
+ Image.fromarray(im_cat).save(os.path.join(cache_path, '%06d.png' % iters))
\ No newline at end of file
diff --git a/legacy/Pretrain/utils/torch_utils.py b/legacy/Pretrain/utils/torch_utils.py
new file mode 100644
index 0000000..b56b616
--- /dev/null
+++ b/legacy/Pretrain/utils/torch_utils.py
@@ -0,0 +1,17 @@
+from distutils.version import LooseVersion, StrictVersion
+import torch
+
+# align_corners option available for v1.3.0 and above
+HAS_AFFINE_ALIGN_CORNERS = LooseVersion(torch.__version__) >= LooseVersion('1.3.0')
+# align_corners defaults to True before v1.4.0, False from v1.4.0 and after
+AFFINE_ALIGN_CORNERS_DEFAULT = LooseVersion(torch.__version__) <= LooseVersion('1.3.0')
+
+
+def affine_align_corners_kw(val):
+ if HAS_AFFINE_ALIGN_CORNERS:
+ return dict(align_corners=val)
+ else:
+ if not val:
+ raise RuntimeError('align_corners not available in torch version {} so '
+ 'cannot set to False'.format(torch.__version__))
+ return {}
diff --git a/legacy/Pretrain/utils/utils.py b/legacy/Pretrain/utils/utils.py
new file mode 100644
index 0000000..6fc3be1
--- /dev/null
+++ b/legacy/Pretrain/utils/utils.py
@@ -0,0 +1,57 @@
+import torch
+import random
+import numpy as np
+import subprocess
+
+def setup_seed(seed):
+ torch.manual_seed(seed)
+ # torch.cuda.manual_seed(seed)
+ torch.cuda.manual_seed_all(seed)
+ np.random.seed(seed)
+ random.seed(seed)
+ torch.backends.cudnn.deterministic = True
+
+def execute(cmd):
+ popen = subprocess.Popen(cmd, stdout=subprocess.PIPE, universal_newlines=True)
+ for stdout_line in iter(popen.stdout.readline, ""):
+ yield stdout_line
+ popen.stdout.close()
+ return_code = popen.wait()
+ if return_code:
+ raise subprocess.CalledProcessError(return_code, cmd)
+
+def compute_num_single(size, stride):
+ # 计算大概需要滑窗的个数
+ num_window = size // stride
+ # 判断是否能整除，如果不可以就+1
+ # 即滑窗个数加1
+ if size % stride != 0:
+ num_window += 1
+ # 计算padding的大小
+ padding_2times = num_window * stride - size + (stride * 2)
+ # 如果padding的大小不能被2整除
+ # 就继续增加滑窗的个数，已使得padding_2times能被2整除
+ while padding_2times % 2 != 0:
+ num_window += 1
+ padding_2times = num_window * stride - size + (stride * 2)
+ # 除以2是为了对称padding
+ padding = padding_2times // 2
+ # 加1是为了补上最后一个完整的滑窗
+ num_window += 1
+ return num_window, padding
+
+def compute_num(raw_shape, stride):
+ size_z = raw_shape[0]
+ size_xy = raw_shape[1]
+ stride_z = stride[0]
+ stride_xy = stride[1]
+ num_z, padding_z = compute_num_single(size_z, stride_z)
+ num_xy, padding_xy = compute_num_single(size_xy, stride_xy)
+ return [num_z, num_xy, num_xy], [padding_z, padding_xy, padding_xy]
+
+if __name__ == "__main__":
+ raw = [500, 4096, 4096]
+ stride = [18, 128, 128]
+ num, padding = compute_num(raw, stride)
+ print(num)
+ print(padding)
diff --git a/legacy/README.md b/legacy/README.md
new file mode 100644
index 0000000..6a857b8
--- /dev/null
+++ b/legacy/README.md
@@ -0,0 +1,107 @@
+
+# SegNeuron 
+Official implementation, datasets and trained models of "SegNeuron: 3D Neuron Instance Segmentation in
+ Any EM Volume with a Generalist Model" （[MICCAI 2024](https://papers.miccai.org/miccai-2024/paper/0518_paper.pdf))
+
+
+
+
+[](https://huggingface.co/datasets/yanchaoz/EMNeuron)
+[](https://colab.research.google.com/github/yanchaoz/SegNeuron/blob/main/SegNeuron_Colab_Inference.ipynb)
+
+
+> [!TIP]
+> **SegNeuron is now available as an Agent Skill in [EM-Skills](https://github.com/yanchaoz/EM-Skills).**
+>
+> **EM-Skills** is a growing collection of reusable Agent Skills for EM analysis, currently including:
+> - 🧠 **Neuron Segmentation Skill** — SegNeuron-based neuron reconstruction
+> - 🧬 **Mitochondria Segmentation Skill** — MitoNet-based mitochondria segmentation
+> - 🎯 **Annotation Selection Skill** — informative region selection for efficient annotation
+> - 🎥 **EM Visualization Skill** — CloudVolume-based visualization and video generation
+>
+> 👉 See **[EM-Skills](https://github.com/yanchaoz/EM-Skills)** for installation and usage. Contributions are welcome—feel free to submit issues, propose new EM Skills, or improve existing ones.
+
+
+## How does SegNeuron speed up neuron segmentation in EM volumes?
+The general-purpose model achieves outstanding reconstruction performance on entirely unseen 3D EM datasets (x/y resolution: **5–10** nm). Human experts only need to perform connectivity corrections on the coarse segmentation results, which can then be directly used to fine-tune SegNeuron or to train new lightweight models.
+
+ 
+
+
+
+
+
+
+## Environments
+We have packaged all the dependencies into Connect.tar.gz, which can be directly downloaded for easy access [here](https://huggingface.co/yanchaoz/SegNeuron).
+## Datasets and Models
+The datasets required for model development and validation are available [here](https://huggingface.co/datasets/yanchaoz/EMNeuron). The trained models can be download [here](https://huggingface.co/yanchaoz/SegNeuron). If you use any of the following vEM datasets in your work, please also cite the corresponding original publications:
+
+- **vEM1: MiRA-ADWT**
+ *Ultrastructural Alterations of Dendritic Morphology in the Prefrontal Cortex of Alzheimer’s Disease Model Rats* [link](https://link.springer.com/article/10.1007/s12264-026-01606-5)
+
+- **vEM2: MiRA-ZF**
+ *Multiplexed Neuromodulatory-Type-Annotated EM-Reconstruction of Larval Zebrafish* [link](https://www.biorxiv.org/content/10.1101/2025.06.12.659365v1)
+
+- **vEM3: MiRA-SCN**
+ *Connectomic Organization of the Suprachiasmatic Nucleus* [link](https://www.biorxiv.org/content/10.1101/2024.10.20.619252v1)
+
+- **vEM4: MiRA-PIB**
+ *PIB: Parallel ion beam etching of sections collected on wafer for ultra large-scale connectomics* [link](https://www.biorxiv.org/content/10.1101/2025.04.25.650569v4)
+
+### Table: Details of EMNeuron
+
+
+
+| Dataset | Modality | Res.($nm$) ($x\/y,z$) | Total voxels (M) | Labeled voxels (M) | Dataset | Modality | Res.($nm$) ($x\/y,z$) | Total voxels (M) | Labeled voxels (M) |
+|----------------------|------------|----------------------|------------------|--------------------|-----------------------|------------|----------------------|------------------|--------------------|
+| ZFinch | SBF-SEM | 9, 20 | 3635 | 131 | HBrain | FIB-SEM | 8, 8 | 3072 | 844 |
+| Layer4 | SBF-SEM | 9, 20 | 1674 | - | FIB25 | FIB-SEM | 8, 8 | 312 | 312 |
+| _vEM1_ (adwt) | ATUM-SEM | 8, 50 | 1205 | 157 | Minnie | ssTEM | 8, 40 | 2096 | - |
+| _vEM2_ (zfish) | ATUM-SEM | 8, 30 | 1329 | 281 | Pinky | ssTEM | 8, 40 | 1165 | 117 |
+| _vEM3_ (scn) | ATUM-SEM | 8, 40 | 1301 | 253 | FAFB | ssTEM | 8, 40 | 2625 | 577 |
+|MitoEM | ATUM-SEM | 8, 30 | 1048 | - | Basil | ssTEM | 8, 40 | 23 | 23 |
+| H01 | ATUM-SEM | 8, 30 | 1166 | 118 | Harris | others | 6, 50 | 30 | 30 |
+| Kasthuri | ATUM-SEM | 6, 30 | 1526 | 478 | _vEM4_ (ionsem) | others | 8, 20 | 45 | - |
+
+
+
+
+
+## Training
+### 1. Pretraining
+```
+cd Pretrain
+```
+```
+python pretrain.py
+```
+### 2. Supervised Training
+```
+cd Train_and_Inference
+```
+```
+python supervised_train.py
+```
+## Inference
+### 1. Affinity Inference
+```
+cd Train_and_Inference
+```
+```
+python inference.py
+```
+### 2. Instance Segmentation
+```
+cd Postprocess
+```
+```
+python FRMC_post.py
+```
+### 3. Zero-shot Segmentation Examples on [MitoEM](https://mitoem.grand-challenge.org/) and [Wildenberg](https://bossdb.org/project/wildenberg2023) (scale bar: 2 um)
+
+ 
+
+
+## Acknowledgement
+This code is based on [SSNS-Net](https://github.com/weih527/SSNS-Net) (IEEE TMI'22) by Huang Wei et al. The postprocessing tools are based on [constantinpape/elf](https://github.com/constantinpape/elf). Should you have any further questions, please let us know. Thanks again for your interest.
diff --git a/legacy/SNAPSHOT.md b/legacy/SNAPSHOT.md
new file mode 100644
index 0000000..1ca26ce
--- /dev/null
+++ b/legacy/SNAPSHOT.md
@@ -0,0 +1,10 @@
+# Original repository snapshot
+
+All files tracked at `yanchaoz/SegNeuron` commit
+`ccb0ba2c5e28e0d2c454e7320c341e71f4eb148c` are preserved here byte for byte,
+including the original README, dependency export, notebook, figures, training
+code and inference scripts. `SNAPSHOT.md` is the only added file in this snapshot.
+
+These files document the original implementation. Use the repository-root README
+for the maintained commands. The historical scripts contain machine-specific
+paths and are not the supported entry points.
diff --git a/legacy/SegNeuron_Colab_Inference.ipynb b/legacy/SegNeuron_Colab_Inference.ipynb
new file mode 100644
index 0000000..61155d9
--- /dev/null
+++ b/legacy/SegNeuron_Colab_Inference.ipynb
@@ -0,0 +1,810 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# SegNeuron Colab: Predict and Segment New EM Volumes\n",
+ "\n",
+ "This is a **single-file, self-contained** Colab project. The model architecture, general-purpose sliding-window inference, blending logic, and FRMC postprocessing are all included in this notebook. It does not clone the repository or import any external Python files.\n",
+ "\n",
+ "**Intended users:** researchers who want to apply a trained SegNeuron model to their own 3D EM TIFF volumes. \n",
+ "**Required inputs:** one `raw.tif/.tiff` volume and the `SegNeuronModel.ckpt` checkpoint. \n",
+ "**Outputs:** affinity predictions, boundary predictions, and an instance segmentation.\n",
+ "\n",
+ "> In Colab, select **Runtime → Change runtime type → T4 GPU**. The checkpoint is about 155 MB, so Google Drive is recommended, although direct browser upload is also supported."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Workflow\n",
+ "\n",
+ "1. Install dependencies and verify GPU availability\n",
+ "2. Upload files from the browser or mount Google Drive\n",
+ "3. Configure input, checkpoint, output, and inference parameters\n",
+ "4. Define the embedded MNet, sliding-window inference, and FRMC postprocessing\n",
+ "5. Run prediction and instance segmentation in one call\n",
+ "6. Inspect and download the results"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Install dependencies that are not included in the default Colab runtime.\n",
+ "# python-elf provides the FRMC/multicut postprocessing used by SegNeuron.\n",
+ "%pip install -q \"python-elf==0.9.1\" \"imageio>=2.31\" \"tifffile>=2023.7\" \"scikit-image>=0.20\" tqdm\n",
+ "\n",
+ "import sys, torch, numpy as np\n",
+ "print(\"Python:\", sys.version.split()[0])\n",
+ "print(\"PyTorch:\", torch.__version__)\n",
+ "print(\"CUDA available:\", torch.cuda.is_available())\n",
+ "if not torch.cuda.is_available():\n",
+ " print(\"WARNING: No GPU is available. 3D inference will be very slow; switch to a GPU runtime.\")"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 1. Prepare the Input Files\n",
+ "\n",
+ "Choose one method:\n",
+ "\n",
+ "- For smaller files, set `UPLOAD_FROM_BROWSER=True`, run the cell, and select both the TIFF and checkpoint.\n",
+ "- For the 155 MB checkpoint and larger TIFF volumes, place the files in Google Drive and set `USE_GOOGLE_DRIVE=True`.\n",
+ "\n",
+ "After uploading or mounting, enter the actual paths in the configuration cell."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "#@title Input source\n",
+ "UPLOAD_FROM_BROWSER = False #@param {type:\"boolean\"}\n",
+ "USE_GOOGLE_DRIVE = False #@param {type:\"boolean\"}\n",
+ "\n",
+ "if UPLOAD_FROM_BROWSER:\n",
+ " from google.colab import files\n",
+ " uploaded = files.upload()\n",
+ " print(\"Uploaded:\", list(uploaded))\n",
+ "\n",
+ "if USE_GOOGLE_DRIVE:\n",
+ " from google.colab import drive\n",
+ " drive.mount(\"/content/drive\")\n",
+ " print(\"Google Drive mounted at /content/drive\")"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "#@title Paths and parameters (edit, then run)\n",
+ "RAW_PATH = \"/content/raw.tif\" #@param {type:\"string\"}\n",
+ "CHECKPOINT_PATH = \"/content/SegNeuronModel.ckpt\" #@param {type:\"string\"}\n",
+ "OUTPUT_DIR = \"/content/segneuron_results\" #@param {type:\"string\"}\n",
+ "\n",
+ "# Values are ordered as z/y/x. The y and x crop sizes should be multiples of 16.\n",
+ "# A stride must not exceed the corresponding crop size.\n",
+ "CROP_SIZE = (20, 128, 128)\n",
+ "STRIDE = (10, 64, 64)\n",
+ "BATCH_SIZE = 1 #@param {type:\"integer\"}\n",
+ "NORMALIZATION = \"auto\" #@param [\"auto\", \"minmax\", \"none\"]\n",
+ "WEIGHT_SIGMA = 0.5 #@param {type:\"number\"}\n",
+ "BETA = 0.25 #@param {type:\"number\"}\n",
+ "\n",
+ "from pathlib import Path\n",
+ "Path(OUTPUT_DIR).mkdir(parents=True, exist_ok=True)\n",
+ "AFFINITY_OUTPUT = str(Path(OUTPUT_DIR) / \"affinities.npy\")\n",
+ "BOUNDARY_OUTPUT = str(Path(OUTPUT_DIR) / \"boundaries.tif\")\n",
+ "SEGMENTATION_OUTPUT = str(Path(OUTPUT_DIR) / \"instances.npy\")\n",
+ "print(\"Results will be written to\", OUTPUT_DIR)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 2. Embedded MNet Model\n",
+ "\n",
+ "The following cell contains the complete MNet definition used by SegNeuron. It is embedded directly in this notebook and does not load code from GitHub or a local `.py` file."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from torch import nn\n",
+ "from torch import nn\n",
+ "import torch\n",
+ "import torch.nn.functional as F\n",
+ "\n",
+ "\n",
+ "class CNA3d(nn.Module): # conv + norm + activation\n",
+ " def __init__(self, in_channels, out_channels, kSize, stride, padding=(1, 1, 1), bias=True, norm_args=None,\n",
+ " activation_args=None):\n",
+ " super().__init__()\n",
+ " self.norm_args = norm_args\n",
+ " self.activation_args = activation_args\n",
+ "\n",
+ " self.conv = nn.Conv3d(in_channels, out_channels, kernel_size=kSize, stride=stride, padding=padding, bias=bias)\n",
+ "\n",
+ " if norm_args is not None:\n",
+ " self.norm = nn.InstanceNorm3d(out_channels, **norm_args)\n",
+ "\n",
+ " if activation_args is not None:\n",
+ " self.activation = nn.LeakyReLU(**activation_args)\n",
+ "\n",
+ " def forward(self, x):\n",
+ " x = self.conv(x)\n",
+ "\n",
+ " if self.norm_args is not None:\n",
+ " x = self.norm(x)\n",
+ "\n",
+ " if self.activation_args is not None:\n",
+ " x = self.activation(x)\n",
+ " return x\n",
+ "\n",
+ "\n",
+ "class CB3d(nn.Module): # conv block 3d\n",
+ " def __init__(self, in_channels, out_channels, kSize=(3, 3), stride=(1, 1), padding=(1, 1, 1), bias=True,\n",
+ " norm_args: tuple = (None, None), activation_args: tuple = (None, None)):\n",
+ " super().__init__()\n",
+ "\n",
+ " self.conv1 = CNA3d(in_channels, out_channels, kSize=kSize[0], stride=stride[0],\n",
+ " padding=padding, bias=bias, norm_args=norm_args[0], activation_args=activation_args[0])\n",
+ "\n",
+ " self.conv2 = CNA3d(out_channels, out_channels, kSize=kSize[1], stride=stride[1],\n",
+ " padding=padding, bias=bias, norm_args=norm_args[1], activation_args=activation_args[1])\n",
+ "\n",
+ " def forward(self, x):\n",
+ " x = self.conv1(x)\n",
+ " x = self.conv2(x)\n",
+ " return x\n",
+ "\n",
+ "\n",
+ "class BasicNet(nn.Module):\n",
+ " norm_kwargs = {'affine': True}\n",
+ " activation_kwargs = {'negative_slope': 1e-2, 'inplace': True}\n",
+ "\n",
+ " def __init__(self):\n",
+ " super(BasicNet, self).__init__()\n",
+ "\n",
+ " def parameter_count(self):\n",
+ " print(\"model have {} paramerters in total\".format(sum(x.numel() for x in self.parameters()) / 1e6))\n",
+ "\n",
+ "\n",
+ "def FMU(x1, x2, mode='sub'):\n",
+ " \"\"\"\n",
+ " feature merging unit\n",
+ " Args:\n",
+ " x1:\n",
+ " x2:\n",
+ " mode: type of fusion\n",
+ " Returns:\n",
+ " \"\"\"\n",
+ " if mode == 'sum':\n",
+ " return torch.add(x1, x2)\n",
+ " elif mode == 'sub':\n",
+ " return torch.abs(x1 - x2)\n",
+ " elif mode == 'cat':\n",
+ " return torch.cat((x1, x2), dim=1)\n",
+ " else:\n",
+ " raise Exception('Unexpected mode')\n",
+ "\n",
+ "\n",
+ "class Down(BasicNet):\n",
+ " def __init__(self, in_channels, out_channels, mode: tuple, FMU='sub', downsample=True, min_z=8):\n",
+ " \"\"\"\n",
+ " basic module at downsampling stage\n",
+ " Args:\n",
+ " in_channels:\n",
+ " out_channels:\n",
+ " mode: represent the streams coming in and out. e.g., ('2d', 'both'): one input stream (2d) and two output streams (2d and 3d)\n",
+ " FMU: determine the type of feature fusion if there are two input streams\n",
+ " downsample: determine whether to downsample input features (only the first module of MNet do not downsample)\n",
+ " min_z: if the size of z-axis < min_z, maxpooling won't be applied along z-axis\n",
+ " \"\"\"\n",
+ " super().__init__()\n",
+ " self.mode_in, self.mode_out = mode\n",
+ " self.downsample = downsample\n",
+ " self.FMU = FMU\n",
+ " self.min_z = min_z\n",
+ " norm_args = (self.norm_kwargs, self.norm_kwargs)\n",
+ " activation_args = (self.activation_kwargs, self.activation_kwargs)\n",
+ "\n",
+ " if self.mode_out == '2d' or self.mode_out == 'both':\n",
+ " self.CB2d = CB3d(in_channels=in_channels, out_channels=out_channels,\n",
+ " kSize=((1, 3, 3), (1, 3, 3)), stride=(1, 1), padding=(0, 1, 1),\n",
+ " norm_args=norm_args, activation_args=activation_args)\n",
+ "\n",
+ " if self.mode_out == '3d' or self.mode_out == 'both':\n",
+ " self.CB3d = CB3d(in_channels=in_channels, out_channels=out_channels,\n",
+ " kSize=(3, 3), stride=(1, 1), padding=(1, 1, 1),\n",
+ " norm_args=norm_args, activation_args=activation_args)\n",
+ "\n",
+ " def forward(self, x):\n",
+ " if self.downsample:\n",
+ " if self.mode_in == 'both':\n",
+ " x2d, x3d = x\n",
+ " p2d = F.max_pool3d(x2d, kernel_size=(1, 2, 2), stride=(1, 2, 2))\n",
+ " if x3d.shape[2] >= self.min_z:\n",
+ " p3d = F.max_pool3d(x3d, kernel_size=(2, 2, 2), stride=(2, 2, 2))\n",
+ " else:\n",
+ " p3d = F.max_pool3d(x3d, kernel_size=(1, 2, 2), stride=(1, 2, 2))\n",
+ "\n",
+ " x = FMU(p2d, p3d, mode=self.FMU)\n",
+ "\n",
+ " elif self.mode_in == '2d':\n",
+ " x = F.max_pool3d(x, kernel_size=(1, 2, 2), stride=(1, 2, 2))\n",
+ "\n",
+ " elif self.mode_in == '3d':\n",
+ " if x.shape[2] >= self.min_z:\n",
+ " x = F.max_pool3d(x, kernel_size=(2, 2, 2), stride=(2, 2, 2))\n",
+ " else:\n",
+ " x = F.max_pool3d(x, kernel_size=(1, 2, 2), stride=(1, 2, 2))\n",
+ "\n",
+ " if self.mode_out == '2d':\n",
+ " return self.CB2d(x)\n",
+ " elif self.mode_out == '3d':\n",
+ " return self.CB3d(x)\n",
+ " elif self.mode_out == 'both':\n",
+ " return self.CB2d(x), self.CB3d(x)\n",
+ "\n",
+ "\n",
+ "class Up(BasicNet):\n",
+ " def __init__(self, in_channels, out_channels, mode: tuple, FMU='sub'):\n",
+ " \"\"\"\n",
+ " basic module at upsampling stage\n",
+ " Args:\n",
+ " in_channels:\n",
+ " out_channels:\n",
+ " mode: represent the streams coming in and out. e.g., ('2d', 'both'): one input stream (2d) and two output streams (2d and 3d)\n",
+ " FMU: determine the type of feature fusion if there are two input streams\n",
+ " \"\"\"\n",
+ " super().__init__()\n",
+ " self.mode_in, self.mode_out = mode\n",
+ " self.FMU = FMU\n",
+ " norm_args = (self.norm_kwargs, self.norm_kwargs)\n",
+ " activation_args = (self.activation_kwargs, self.activation_kwargs)\n",
+ "\n",
+ " if self.mode_out == '2d' or self.mode_out == 'both':\n",
+ " self.CB2d = CB3d(in_channels=in_channels, out_channels=out_channels,\n",
+ " kSize=((1, 3, 3), (1, 3, 3)), stride=(1, 1), padding=(0, 1, 1),\n",
+ " norm_args=norm_args, activation_args=activation_args)\n",
+ "\n",
+ " if self.mode_out == '3d' or self.mode_out == 'both':\n",
+ " self.CB3d = CB3d(in_channels=in_channels, out_channels=out_channels,\n",
+ " kSize=(3, 3), stride=(1, 1), padding=(1, 1, 1),\n",
+ " norm_args=norm_args, activation_args=activation_args)\n",
+ "\n",
+ " def forward(self, x):\n",
+ " x2d, xskip2d, x3d, xskip3d = x\n",
+ "\n",
+ " tarSize = xskip2d.shape[2:]\n",
+ " up2d = F.interpolate(x2d, size=tarSize, mode='trilinear', align_corners=False)\n",
+ " up3d = F.interpolate(x3d, size=tarSize, mode='trilinear', align_corners=False)\n",
+ "\n",
+ " cat = torch.cat([FMU(xskip2d, xskip3d, self.FMU), FMU(up2d, up3d, self.FMU)], dim=1)\n",
+ "\n",
+ " if self.mode_out == '2d':\n",
+ " return self.CB2d(cat)\n",
+ " elif self.mode_out == '3d':\n",
+ " return self.CB3d(cat)\n",
+ " elif self.mode_out == 'both':\n",
+ " return self.CB2d(cat), self.CB3d(cat)\n",
+ "\n",
+ "\n",
+ "class MNet(BasicNet):\n",
+ " def __init__(self, in_channels, kn=(32, 48, 64, 80, 96), FMU='sub'):\n",
+ " \"\"\"\n",
+ "\n",
+ " Args:\n",
+ " in_channels: channels of input\n",
+ " num_classes: output classes\n",
+ " kn: the number of kernels\n",
+ " ds: deep supervision\n",
+ " FMU: type of feature merging unit\n",
+ " \"\"\"\n",
+ " super().__init__()\n",
+ "\n",
+ " channel_factor = {'sum': 1, 'sub': 1, 'cat': 2}\n",
+ " fct = channel_factor[FMU]\n",
+ "\n",
+ " self.down11 = Down(in_channels, kn[0], ('/', 'both'), downsample=False)\n",
+ " self.down12 = Down(kn[0], kn[1], ('2d', 'both'))\n",
+ " self.down13 = Down(kn[1], kn[2], ('2d', 'both'))\n",
+ " self.down14 = Down(kn[2], kn[3], ('2d', 'both'))\n",
+ " self.bottleneck1 = Down(kn[3], kn[4], ('2d', '2d'))\n",
+ " self.up11 = Up(fct * (kn[3] + kn[4]), kn[3], ('both', '2d'), FMU)\n",
+ " self.up12 = Up(fct * (kn[2] + kn[3]), kn[2], ('both', '2d'), FMU)\n",
+ " self.up13 = Up(fct * (kn[1] + kn[2]), kn[1], ('both', '2d'), FMU)\n",
+ " self.up14 = Up(fct * (kn[0] + kn[1]), kn[0], ('both', 'both'), FMU)\n",
+ "\n",
+ " self.up11_ = Up(fct * (kn[3] + kn[4]), kn[3], ('both', '2d'), FMU)\n",
+ " self.up12_ = Up(fct * (kn[2] + kn[3]), kn[2], ('both', '2d'), FMU)\n",
+ " self.up13_ = Up(fct * (kn[1] + kn[2]), kn[1], ('both', '2d'), FMU)\n",
+ " self.up14_ = Up(fct * (kn[0] + kn[1]), kn[0], ('both', 'both'), FMU)\n",
+ "\n",
+ " self.down21 = Down(kn[0], kn[1], ('3d', 'both'))\n",
+ " self.down22 = Down(fct * kn[1], kn[2], ('both', 'both'), FMU)\n",
+ " self.down23 = Down(fct * kn[2], kn[3], ('both', 'both'), FMU)\n",
+ " self.bottleneck2 = Down(fct * kn[3], kn[4], ('both', 'both'), FMU)\n",
+ " self.up21 = Up(fct * (kn[3] + kn[4]), kn[3], ('both', 'both'), FMU)\n",
+ " self.up22 = Up(fct * (kn[2] + kn[3]), kn[2], ('both', 'both'), FMU)\n",
+ " self.up23 = Up(fct * (kn[1] + kn[2]), kn[1], ('both', '3d'), FMU)\n",
+ "\n",
+ " self.up21_ = Up(fct * (kn[3] + kn[4]), kn[3], ('both', 'both'), FMU)\n",
+ " self.up22_ = Up(fct * (kn[2] + kn[3]), kn[2], ('both', 'both'), FMU)\n",
+ " self.up23_ = Up(fct * (kn[1] + kn[2]), kn[1], ('both', '3d'), FMU)\n",
+ "\n",
+ " self.down31 = Down(kn[1], kn[2], ('3d', 'both'))\n",
+ " self.down32 = Down(fct * kn[2], kn[3], ('both', 'both'), FMU)\n",
+ " self.bottleneck3 = Down(fct * kn[3], kn[4], ('both', 'both'), FMU)\n",
+ " self.up31 = Up(fct * (kn[3] + kn[4]), kn[3], ('both', 'both'), FMU)\n",
+ " self.up32 = Up(fct * (kn[2] + kn[3]), kn[2], ('both', '3d'), FMU)\n",
+ "\n",
+ " self.up31_ = Up(fct * (kn[3] + kn[4]), kn[3], ('both', 'both'), FMU)\n",
+ " self.up32_ = Up(fct * (kn[2] + kn[3]), kn[2], ('both', '3d'), FMU)\n",
+ "\n",
+ " self.down41 = Down(kn[2], kn[3], ('3d', 'both'), FMU)\n",
+ " self.bottleneck4 = Down(fct * kn[3], kn[4], ('both', 'both'), FMU)\n",
+ " self.up41 = Up(fct * (kn[3] + kn[4]), kn[3], ('both', '3d'), FMU)\n",
+ "\n",
+ " self.up41_ = Up(fct * (kn[3] + kn[4]), kn[3], ('both', '3d'), FMU)\n",
+ "\n",
+ " self.bottleneck5 = Down(kn[3], kn[4], ('3d', '3d'))\n",
+ "\n",
+ " self.outputs = nn.ModuleList(\n",
+ " [nn.Conv3d(c, 3, kernel_size=(1, 1, 1), stride=1, padding=0, bias=False)\n",
+ " for c in [kn[0], kn[1], kn[1], kn[2], kn[2], kn[3], kn[3]]]\n",
+ " )\n",
+ "\n",
+ " self.outputs2 = nn.ModuleList(\n",
+ " [nn.Conv3d(c, 1, kernel_size=(1, 1, 1), stride=1, padding=0, bias=False)\n",
+ " for c in [kn[0], kn[1], kn[1], kn[2], kn[2], kn[3], kn[3]]]\n",
+ " )\n",
+ " self.sigmoid = nn.Sigmoid()\n",
+ "\n",
+ " def forward(self, x):\n",
+ " down11 = self.down11(x)\n",
+ " down12 = self.down12(down11[0])\n",
+ " down13 = self.down13(down12[0])\n",
+ " down14 = self.down14(down13[0])\n",
+ " bottleNeck1 = self.bottleneck1(down14[0])\n",
+ "\n",
+ " down21 = self.down21(down11[1])\n",
+ " down22 = self.down22([down21[0], down12[1]])\n",
+ " down23 = self.down23([down22[0], down13[1]])\n",
+ " bottleNeck2 = self.bottleneck2([down23[0], down14[1]])\n",
+ "\n",
+ " down31 = self.down31(down21[1])\n",
+ " down32 = self.down32([down31[0], down22[1]])\n",
+ " bottleNeck3 = self.bottleneck3([down32[0], down23[1]])\n",
+ "\n",
+ " down41 = self.down41(down31[1])\n",
+ " bottleNeck4 = self.bottleneck4([down41[0], down32[1]])\n",
+ "\n",
+ " bottleNeck5 = self.bottleneck5(down41[1])\n",
+ "\n",
+ " up41 = self.up41([bottleNeck4[0], down41[0], bottleNeck5, down41[1]])\n",
+ "\n",
+ " up31 = self.up31([bottleNeck3[0], down32[0], bottleNeck4[1], down32[1]])\n",
+ " up32 = self.up32([up31[0], down31[0], up41, down31[1]])\n",
+ "\n",
+ " up21 = self.up21([bottleNeck2[0], down23[0], bottleNeck3[1], down23[1]])\n",
+ " up22 = self.up22([up21[0], down22[0], up31[1], down22[1]])\n",
+ " up23 = self.up23([up22[0], down21[0], up32, down21[1]])\n",
+ "\n",
+ " up11 = self.up11([bottleNeck1, down14[0], bottleNeck2[1], down14[1]])\n",
+ " up12 = self.up12([up11, down13[0], up21[1], down13[1]])\n",
+ " up13 = self.up13([up12, down12[0], up22[1], down12[1]])\n",
+ " up14 = self.up14([up13, down11[0], up23, down11[1]])\n",
+ "\n",
+ " up41_ = self.up41_([bottleNeck4[0], down41[0], bottleNeck5, down41[1]])\n",
+ "\n",
+ " up31_ = self.up31_([bottleNeck3[0], down32[0], bottleNeck4[1], down32[1]])\n",
+ " up32_ = self.up32_([up31_[0], down31[0], up41_, down31[1]])\n",
+ "\n",
+ " up21_ = self.up21_([bottleNeck2[0], down23[0], bottleNeck3[1], down23[1]])\n",
+ " up22_ = self.up22_([up21_[0], down22[0], up31_[1], down22[1]])\n",
+ " up23_ = self.up23_([up22_[0], down21[0], up32_, down21[1]])\n",
+ "\n",
+ " up11_ = self.up11_([bottleNeck1, down14[0], bottleNeck2[1], down14[1]])\n",
+ " up12_ = self.up12_([up11_, down13[0], up21_[1], down13[1]])\n",
+ " up13_ = self.up13_([up12_, down12[0], up22_[1], down12[1]])\n",
+ " up14_ = self.up14_([up13_, down11[0], up23_, down11[1]])\n",
+ "\n",
+ " return self.sigmoid(self.outputs[0](up14[0] + up14[1])), self.sigmoid(self.outputs2[0](up14_[0] + up14_[1]))"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 3. General-Purpose Sliding-Window Inference\n",
+ "\n",
+ "The implementation supports rectangular and non-divisible volumes, pads axes smaller than the crop, places the final window flush with each boundary, guarantees full coverage, blends overlaps with Gaussian weights, handles batches safely, and recognizes common checkpoint dictionary formats."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from collections import OrderedDict\n",
+ "from pathlib import Path\n",
+ "import math\n",
+ "import imageio.v2 as imageio\n",
+ "import numpy as np\n",
+ "import tifffile\n",
+ "import torch\n",
+ "from torch.utils.data import Dataset, DataLoader\n",
+ "from tqdm.auto import tqdm\n",
+ "\n",
+ "\n",
+ "def _triple(value, name):\n",
+ " if len(value) != 3 or any(int(v) <= 0 for v in value):\n",
+ " raise ValueError(f\"{name} must contain three positive integers (z, y, x): {value}\")\n",
+ " return tuple(int(v) for v in value)\n",
+ "\n",
+ "\n",
+ "def window_starts(length, window, step):\n",
+ " if step > window:\n",
+ " raise ValueError(f\"stride {step} cannot exceed crop size {window}\")\n",
+ " if length <= window:\n",
+ " return [0]\n",
+ " starts = list(range(0, length - window + 1, step))\n",
+ " if starts[-1] != length - window:\n",
+ " starts.append(length - window)\n",
+ " return starts\n",
+ "\n",
+ "\n",
+ "def normalize_volume(raw, mode=\"auto\"):\n",
+ " raw = np.asarray(raw)\n",
+ " if mode == \"none\":\n",
+ " return raw.astype(np.float32, copy=False)\n",
+ " if mode == \"minmax\":\n",
+ " out = raw.astype(np.float32)\n",
+ " lo, hi = float(np.nanmin(out)), float(np.nanmax(out))\n",
+ " return (out - lo) / (hi - lo) if hi > lo else np.zeros_like(out)\n",
+ " if mode != \"auto\":\n",
+ " raise ValueError(f\"unknown normalization: {mode}\")\n",
+ " if np.issubdtype(raw.dtype, np.integer):\n",
+ " return raw.astype(np.float32) / float(np.iinfo(raw.dtype).max)\n",
+ " out = raw.astype(np.float32, copy=False)\n",
+ " lo, hi = float(np.nanmin(out)), float(np.nanmax(out))\n",
+ " if 0 <= lo and hi <= 1:\n",
+ " return out\n",
+ " return (out - lo) / (hi - lo) if hi > lo else np.zeros_like(out)\n",
+ "\n",
+ "\n",
+ "class SlidingWindowVolume(Dataset):\n",
+ " def __init__(self, raw_path, crop_size=(20, 128, 128), stride=(10, 64, 64),\n",
+ " normalization=\"auto\", weight_sigma=0.5):\n",
+ " raw_path = Path(raw_path).expanduser()\n",
+ " if not raw_path.is_file():\n",
+ " raise FileNotFoundError(f\"Input TIFF not found: {raw_path}\")\n",
+ " raw = np.asarray(tifffile.imread(raw_path))\n",
+ " if raw.ndim != 3:\n",
+ " raise ValueError(f\"Expected TIFF shape (z,y,x), got {raw.shape}\")\n",
+ " self.original_shape = tuple(raw.shape)\n",
+ " self.crop_size = _triple(crop_size, \"crop_size\")\n",
+ " self.stride = _triple(stride, \"stride\")\n",
+ " if self.crop_size[1] % 16 or self.crop_size[2] % 16:\n",
+ " raise ValueError(\"crop y/x must be multiples of 16 for MNet\")\n",
+ " pad_after = tuple(max(crop - size, 0) for size, crop in zip(raw.shape, self.crop_size))\n",
+ " if any(pad_after):\n",
+ " raw = np.pad(raw, tuple((0, p) for p in pad_after), mode=\"edge\")\n",
+ " self.raw = normalize_volume(raw, normalization)\n",
+ " self.padded_shape = tuple(self.raw.shape)\n",
+ " axes = [window_starts(n, c, s) for n, c, s in\n",
+ " zip(self.padded_shape, self.crop_size, self.stride)]\n",
+ " self.positions = [(z, y, x) for z in axes[0] for y in axes[1] for x in axes[2]]\n",
+ " coords = [np.linspace(-1, 1, n, dtype=np.float32) for n in self.crop_size]\n",
+ " zz, yy, xx = np.meshgrid(*coords, indexing=\"ij\")\n",
+ " if weight_sigma > 0:\n",
+ " weight = np.exp(-(zz*zz + yy*yy + xx*xx) / (2 * weight_sigma**2))\n",
+ " weight = np.maximum(weight, 1e-3)\n",
+ " else:\n",
+ " weight = np.ones(self.crop_size, np.float32)\n",
+ " self.weight = weight[None].astype(np.float32)\n",
+ "\n",
+ " def __len__(self): return len(self.positions)\n",
+ "\n",
+ " def __getitem__(self, index):\n",
+ " z, y, x = self.positions[index]\n",
+ " dz, dy, dx = self.crop_size\n",
+ " patch = self.raw[z:z+dz, y:y+dy, x:x+dx]\n",
+ " return np.ascontiguousarray(patch[None], dtype=np.float32), index\n",
+ "\n",
+ " def accumulator(self, channels):\n",
+ " return (np.zeros((channels,) + self.padded_shape, np.float32),\n",
+ " np.zeros((1,) + self.padded_shape, np.float32))\n",
+ "\n",
+ " def add(self, total, weights, prediction, index):\n",
+ " z, y, x = self.positions[int(index)]\n",
+ " dz, dy, dx = self.crop_size\n",
+ " if tuple(prediction.shape[1:]) != self.crop_size:\n",
+ " raise ValueError(f\"Model output {prediction.shape[1:]} != crop {self.crop_size}\")\n",
+ " total[:, z:z+dz, y:y+dy, x:x+dx] += prediction * self.weight\n",
+ " weights[:, z:z+dz, y:y+dy, x:x+dx] += self.weight\n",
+ "\n",
+ " def finish(self, total, weights):\n",
+ " if np.any(weights == 0):\n",
+ " raise RuntimeError(\"Sliding windows left uncovered voxels\")\n",
+ " z, y, x = self.original_shape\n",
+ " return (total / weights)[:, :z, :y, :x]\n",
+ "\n",
+ "\n",
+ "def checkpoint_state_dict(checkpoint):\n",
+ " if isinstance(checkpoint, dict):\n",
+ " for key in (\"model_weights\", \"model_state_dict\", \"state_dict\", \"model\"):\n",
+ " if key in checkpoint and isinstance(checkpoint[key], dict):\n",
+ " checkpoint = checkpoint[key]\n",
+ " break\n",
+ " if not isinstance(checkpoint, dict):\n",
+ " raise ValueError(\"Checkpoint does not contain a state dictionary\")\n",
+ " return OrderedDict((key.removeprefix(\"module.\"), value) for key, value in checkpoint.items())\n",
+ "\n",
+ "\n",
+ "def load_model(checkpoint_path, device):\n",
+ " checkpoint_path = Path(checkpoint_path).expanduser()\n",
+ " if not checkpoint_path.is_file():\n",
+ " raise FileNotFoundError(f\"Checkpoint not found: {checkpoint_path}\")\n",
+ " model = MNet(1, kn=(32, 64, 96, 128, 256), FMU=\"sub\")\n",
+ " state = torch.load(checkpoint_path, map_location=device)\n",
+ " model.load_state_dict(checkpoint_state_dict(state))\n",
+ " return model.to(device).eval()\n",
+ "\n",
+ "\n",
+ "@torch.inference_mode()\n",
+ "def predict(raw_path, checkpoint_path, affinity_output, boundary_output,\n",
+ " crop_size=(20,128,128), stride=(10,64,64), batch_size=1,\n",
+ " normalization=\"auto\", weight_sigma=0.5, device=None):\n",
+ " device = torch.device(device or (\"cuda\" if torch.cuda.is_available() else \"cpu\"))\n",
+ " dataset = SlidingWindowVolume(raw_path, crop_size, stride, normalization, weight_sigma)\n",
+ " loader = DataLoader(dataset, batch_size=batch_size, shuffle=False,\n",
+ " num_workers=0, pin_memory=device.type == \"cuda\")\n",
+ " model = load_model(checkpoint_path, device)\n",
+ " affinity_sum = affinity_weight = boundary_sum = boundary_weight = None\n",
+ " for inputs, indices in tqdm(loader, desc=\"SegNeuron inference\"):\n",
+ " affinities, boundaries = model(inputs.to(device, non_blocking=True))\n",
+ " affinities, boundaries = affinities.cpu().numpy(), boundaries.cpu().numpy()\n",
+ " if affinity_sum is None:\n",
+ " affinity_sum, affinity_weight = dataset.accumulator(affinities.shape[1])\n",
+ " boundary_sum, boundary_weight = dataset.accumulator(boundaries.shape[1])\n",
+ " for affinity, boundary, index in zip(affinities, boundaries, indices.tolist()):\n",
+ " dataset.add(affinity_sum, affinity_weight, affinity, index)\n",
+ " dataset.add(boundary_sum, boundary_weight, boundary, index)\n",
+ " affinity = dataset.finish(affinity_sum, affinity_weight)\n",
+ " boundary = dataset.finish(boundary_sum, boundary_weight)\n",
+ " np.save(affinity_output, affinity)\n",
+ " tifffile.imwrite(boundary_output, boundary[0], photometric=\"minisblack\")\n",
+ " print(\"Affinity:\", affinity.shape, affinity.dtype, \"->\", affinity_output)\n",
+ " print(\"Boundary:\", boundary.shape, boundary.dtype, \"->\", boundary_output)\n",
+ " return affinity, boundary\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 4. Embedded FRMC/Multicut Postprocessing\n",
+ "\n",
+ "This function receives in-memory prediction arrays directly. It implements the watershed, region adjacency graph, and Kernighan–Lin multicut workflow used by the original SegNeuron postprocessing script."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def postprocess(affinities, boundaries=None, output=None, beta=0.25):\n",
+ " import elf.segmentation.features as feats\n",
+ " import elf.segmentation.multicut as mc\n",
+ " import elf.segmentation.watershed as ws\n",
+ "\n",
+ " affinities = np.asarray(affinities, dtype=np.float32)\n",
+ " if affinities.ndim != 4 or affinities.shape[0] != 3:\n",
+ " raise ValueError(f\"Affinities must be (3,z,y,x), got {affinities.shape}\")\n",
+ " if boundaries is not None:\n",
+ " boundaries = np.asarray(boundaries, dtype=np.float32)\n",
+ " if boundaries.ndim == 4 and boundaries.shape[0] == 1:\n",
+ " boundaries = boundaries[0]\n",
+ " if boundaries.shape != affinities.shape[1:]:\n",
+ " raise ValueError(f\"Boundary {boundaries.shape} != {affinities.shape[1:]}\")\n",
+ " affinities = np.minimum(affinities, boundaries[None])\n",
+ "\n",
+ " inverted = 1.0 - affinities\n",
+ " boundary_input = np.maximum(inverted[1], inverted[2])\n",
+ " watershed = np.zeros_like(boundary_input, dtype=np.uint64)\n",
+ " offset = 0\n",
+ " for z in tqdm(range(watershed.shape[0]), desc=\"Watershed\"):\n",
+ " wsz, max_id = ws.distance_transform_watershed(\n",
+ " boundary_input[z], threshold=0.25, sigma_seeds=2.0)\n",
+ " wsz = wsz.astype(np.uint64, copy=False)\n",
+ " wsz[wsz != 0] += offset\n",
+ " offset += int(max_id)\n",
+ " watershed[z] = wsz\n",
+ "\n",
+ " rag = feats.compute_rag(watershed)\n",
+ " offsets = [[-1,0,0], [0,-1,0], [0,0,-1]]\n",
+ " probabilities = feats.compute_affinity_features(rag, inverted, offsets)[:, 0]\n",
+ " edge_sizes = feats.compute_boundary_mean_and_length(rag, boundary_input)[:, 1]\n",
+ " costs = mc.transform_probabilities_to_costs(\n",
+ " probabilities, edge_sizes=edge_sizes, beta=float(beta))\n",
+ " node_labels = mc.multicut_kernighan_lin(rag, costs)\n",
+ " segmentation = feats.project_node_labels_to_pixels(rag, node_labels)\n",
+ " if output:\n",
+ " np.save(output, segmentation)\n",
+ " print(\"Instances:\", segmentation.shape, segmentation.dtype, \"->\", output)\n",
+ " return segmentation\n",
+ "\n",
+ "\n",
+ "def run_segmentation(raw_path, checkpoint_path, output_dir,\n",
+ " crop_size=(20,128,128), stride=(10,64,64), batch_size=1,\n",
+ " normalization=\"auto\", weight_sigma=0.5, beta=0.25):\n",
+ " output_dir = Path(output_dir)\n",
+ " output_dir.mkdir(parents=True, exist_ok=True)\n",
+ " affinities, boundaries = predict(\n",
+ " raw_path, checkpoint_path,\n",
+ " str(output_dir / \"affinities.npy\"), str(output_dir / \"boundaries.tif\"),\n",
+ " crop_size, stride, batch_size, normalization, weight_sigma)\n",
+ " instances = postprocess(\n",
+ " affinities, boundaries, str(output_dir / \"instances.npy\"), beta)\n",
+ " return affinities, boundaries, instances\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 5. Preflight Check\n",
+ "\n",
+ "This cell only reads file metadata and validates paths; it does not start model inference. Confirm that the TIFF shape is displayed as `(z, y, x)`. If the exported volume is `(y, x, z)`, transpose it before inference."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import os, tifffile\n",
+ "\n",
+ "for name, path in {\"raw\": RAW_PATH, \"checkpoint\": CHECKPOINT_PATH}.items():\n",
+ " if not Path(path).is_file():\n",
+ " raise FileNotFoundError(f\"{name} file not found: {path}\")\n",
+ "raw_info = tifffile.imread(RAW_PATH)\n",
+ "if raw_info.ndim != 3:\n",
+ " raise ValueError(f\"RAW must be 3-D, got {raw_info.shape}\")\n",
+ "estimated = (len(window_starts(max(raw_info.shape[0], CROP_SIZE[0]), CROP_SIZE[0], STRIDE[0])) *\n",
+ " len(window_starts(max(raw_info.shape[1], CROP_SIZE[1]), CROP_SIZE[1], STRIDE[1])) *\n",
+ " len(window_starts(max(raw_info.shape[2], CROP_SIZE[2]), CROP_SIZE[2], STRIDE[2])))\n",
+ "print(\"RAW:\", raw_info.shape, raw_info.dtype,\n",
+ " \"range\", (float(np.nanmin(raw_info)), float(np.nanmax(raw_info))))\n",
+ "print(\"Checkpoint:\", round(Path(CHECKPOINT_PATH).stat().st_size / 1024**2, 1), \"MiB\")\n",
+ "print(\"Sliding windows:\", estimated)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 6. Run Prediction and Instance Segmentation\n",
+ "\n",
+ "This cell performs full GPU inference followed by FRMC postprocessing. The in-memory affinity array requires approximately `3 × z × y × x × 4` bytes."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "affinities, boundaries, instances = run_segmentation(\n",
+ " RAW_PATH,\n",
+ " CHECKPOINT_PATH,\n",
+ " OUTPUT_DIR,\n",
+ " crop_size=CROP_SIZE,\n",
+ " stride=STRIDE,\n",
+ " batch_size=BATCH_SIZE,\n",
+ " normalization=NORMALIZATION,\n",
+ " weight_sigma=WEIGHT_SIGMA,\n",
+ " beta=BETA,\n",
+ ")\n",
+ "print(\"Unique instance labels:\", len(np.unique(instances)))"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 7. Inspect the Center Slice"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import matplotlib.pyplot as plt\n",
+ "\n",
+ "z = instances.shape[0] // 2\n",
+ "raw = tifffile.imread(RAW_PATH)\n",
+ "fig, axes = plt.subplots(1, 3, figsize=(18, 6))\n",
+ "axes[0].imshow(raw[z], cmap=\"gray\")\n",
+ "axes[0].set_title(f\"Raw z={z}\")\n",
+ "axes[1].imshow(boundaries[0, z], cmap=\"magma\", vmin=0, vmax=1)\n",
+ "axes[1].set_title(\"Boundary\")\n",
+ "axes[2].imshow(instances[z], cmap=\"nipy_spectral\")\n",
+ "axes[2].set_title(\"Instances\")\n",
+ "for ax in axes: ax.axis(\"off\")\n",
+ "plt.tight_layout()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## 8. Package and Download the Results\n",
+ "\n",
+ "The result directory contains:\n",
+ "\n",
+ "- `affinities.npy`: `(3,z,y,x)` float32\n",
+ "- `boundaries.tif`: `(z,y,x)` float32\n",
+ "- `instances.npy`: `(z,y,x)` uint64"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "#@title Download ZIP\n",
+ "DOWNLOAD_RESULTS = True #@param {type:\"boolean\"}\n",
+ "\n",
+ "import shutil\n",
+ "archive = shutil.make_archive(\"/content/segneuron_results\", \"zip\", OUTPUT_DIR)\n",
+ "print(\"Created\", archive, round(Path(archive).stat().st_size / 1024**2, 1), \"MiB\")\n",
+ "if DOWNLOAD_RESULTS:\n",
+ " from google.colab import files\n",
+ " files.download(archive)"
+ ]
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "Python 3",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "name": "python",
+ "version": "3.12"
+ },
+ "colab": {
+ "name": "SegNeuron_Colab_Inference.ipynb",
+ "provenance": []
+ },
+ "accelerator": "GPU"
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
\ No newline at end of file
diff --git a/legacy/Train_and_Inference/config/SegNeuron.yaml b/legacy/Train_and_Inference/config/SegNeuron.yaml
new file mode 100644
index 0000000..eeef62c
--- /dev/null
+++ b/legacy/Train_and_Inference/config/SegNeuron.yaml
@@ -0,0 +1,50 @@
+NAME: 'SegNeuron'
+
+MODEL:
+ pre_train: True
+ pretrain_path: '/***/***'
+ continue_train: False
+ continue_path: '/***/***'
+
+TRAIN:
+ resume: False
+ if_valid: True
+ cache_path: './caches/'
+ save_path: './models/'
+ pad: 0
+ loss_func: 'BCELoss'
+ opt_type: 'adam'
+ display_freq: 100
+ total_iters: 400000
+ warmup_iters: 0
+ base_lr: 0.01
+ end_lr: 0.0001
+ save_freq: 2000
+ valid_freq: 2000
+ decay_iters: 200000
+ weight_decay: ~
+ power: 1.5
+ batch_size: 8
+ num_workers: 32
+ if_cuda: True
+ random_seed: 666
+ min_valid_iter: 10000
+ freq_mix_prob: 0.25
+ spa_mix_prob: 0.25
+
+DATA:
+ min_noise_std: 0.01
+ max_noise_std: 0.2
+ min_kernel_size: 3
+ max_kernel_size: 9
+ min_sigma: 0
+ max_sigma: 2
+ data_folder: '/***/***'
+ data_folder_val: '/***/***'
+ start_slice: 0
+ end_slice: 100
+ val_start: 0
+ val_end: 100
+ predict_split: False
+ if_ignore_bg: True
+
diff --git a/legacy/Train_and_Inference/inference.py b/legacy/Train_and_Inference/inference.py
new file mode 100644
index 0000000..7aa21b6
--- /dev/null
+++ b/legacy/Train_and_Inference/inference.py
@@ -0,0 +1,67 @@
+import os
+import yaml
+import argparse
+import imageio
+import numpy as np
+from attrdict import AttrDict
+from collections import OrderedDict
+from tqdm import tqdm
+import warnings
+import torch
+import torch.nn as nn
+from inference_provider import Provider_valid
+from model.Mnet import MNet
+
+os.environ['CUDA_VISIBLE_DEVICES'] = "0"
+warnings.filterwarnings("ignore")
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument('-c', '--cfg', type=str, default='SegNeuron', help='path to config file')
+ args = parser.parse_args()
+ cfg_file = args.cfg + '.yaml'
+ print('cfg_file: ' + cfg_file)
+ with open('./config/' + cfg_file, 'r') as f:
+ cfg = AttrDict(yaml.safe_load(f))
+
+ pth = '/***/***.pth'
+
+ model = MNet(1, kn=(32, 64, 96, 128, 256), FMU='sub').cuda()
+ checkpoint = torch.load(pth)
+ new_state_dict = OrderedDict()
+ state_dict = checkpoint['model_weights']
+ for k, v in state_dict.items():
+ name = k.replace('module.', '') if 'module' in k else k
+ new_state_dict[name] = v
+ print('load mnet!')
+ model.load_state_dict(new_state_dict)
+ model = model.cuda()
+
+ model.eval()
+ valid_provider = Provider_valid(cfg, valid_data='***')
+ criterion = nn.BCELoss()
+ dataloader = torch.utils.data.DataLoader(valid_provider, batch_size=1, num_workers=0,
+ shuffle=False, drop_last=False, pin_memory=True)
+
+ pbar = tqdm(total=len(valid_provider))
+ losses_valid = []
+ for k, batch in enumerate(dataloader, 0):
+ inputs, target, _ = batch
+ inputs = inputs.cuda()
+ target = target.cuda()
+ with torch.no_grad():
+ pred, bound = model(inputs)
+ valid_provider.add_vol(np.squeeze(pred.data.cpu().numpy()))
+ valid_provider.add_bound(np.squeeze(bound.data.cpu().numpy()))
+ pbar.update(1)
+ pbar.close()
+
+ out_affs = valid_provider.get_results()
+ out_bounds = valid_provider.get_results_bound()
+
+ gt_affs = valid_provider.get_gt_affs()
+ gt_seg = valid_provider.get_gt_lb()
+ valid_provider.reset_output()
+
+ np.save('/***/***', out_affs)
+ imageio.volwrite('/***/***.tif', out_bounds.squeeze())
diff --git a/legacy/Train_and_Inference/inference_provider.py b/legacy/Train_and_Inference/inference_provider.py
new file mode 100644
index 0000000..0026b5c
--- /dev/null
+++ b/legacy/Train_and_Inference/inference_provider.py
@@ -0,0 +1,201 @@
+import os
+import cv2
+import h5py
+import math
+import random
+import numpy as np
+from PIL import Image
+from torch.utils.data import Dataset
+from utils.seg_util import mknhood3d, genSegMalis
+from utils.aff_util import seg_to_affgraph
+import imageio
+
+
+class Provider_valid(Dataset):
+ def __init__(self, cfg, valid_data=None):
+ # basic settings
+ self.cfg = cfg
+ if valid_data is not None:
+ valid_dataset_name = valid_data
+ else:
+ valid_dataset_name = cfg.DATA.valid_dataset
+ print('valid dataset:', valid_dataset_name)
+
+ self.crop_size = [20, 128, 128]
+ self.stride = [10, 64, 64]
+ self.out_size = [20, 128, 128]
+
+ if valid_dataset_name == 'Harris':
+ self.sub_path = 'OutofDistribution/Harris'
+ self.train_datasets = ['raw.tif']
+ self.train_labels = ['label.tif']
+ elif valid_dataset_name == 'Microns[B]':
+ self.sub_path = 'OutofDistribution/Microns[B]'
+ self.train_datasets = ['raw.tif']
+ self.train_labels = ['label.tif']
+ elif valid_dataset_name == 'IONSEM':
+ self.sub_path = 'OutofDistribution/Mira-ionsem'
+ self.train_datasets = ['raw.tif']
+ self.train_labels = ['label.tif']
+ else:
+ raise AttributeError('No this dataset type!')
+
+ self.folder_name = os.path.join(cfg.DATA.data_folder_val, self.sub_path)
+ assert len(self.train_datasets) == len(self.train_labels)
+
+ # load dataset
+ self.dataset = []
+ self.labels = []
+ for k in range(len(self.train_datasets)):
+ print('load ' + self.folder_name + self.train_datasets[k] + ' ...')
+ data = imageio.volread(os.path.join(self.folder_name, self.train_datasets[k]))
+ self.dataset.append(data[:])
+ label = imageio.volread(os.path.join(self.folder_name, self.train_labels[k]))
+ self.labels.append(label[:])
+ self.origin_data_shape = list(self.dataset[0].shape)
+
+ self.gt_affs = []
+ for k in range(len(self.labels)):
+ temp = self.labels[k].copy()
+ self.gt_affs.append(seg_to_affgraph(temp, mknhood3d(1), pad='replicate').astype(np.float32))
+
+ self.num_zyx = [(self.origin_data_shape[0] - self.out_size[0]) // self.stride[0] + 2,
+ (self.origin_data_shape[1] - self.out_size[1]) // self.stride[1] + 2,
+ (self.origin_data_shape[2] - self.out_size[2]) // self.stride[2] + 2]
+ self.valid_padding = [
+ ((self.num_zyx[0] - 1) * self.stride[0] + self.out_size[0] - self.origin_data_shape[0]) // 2,
+ ((self.num_zyx[1] - 1) * self.stride[1] + self.out_size[1] - self.origin_data_shape[1]) // 2,
+ ((self.num_zyx[2] - 1) * self.stride[2] + self.out_size[2] - self.origin_data_shape[2]) // 2]
+
+ for k in range(len(self.dataset)):
+ self.dataset[k] = np.pad(self.dataset[k], ((self.valid_padding[0], self.valid_padding[0]), \
+ (self.valid_padding[1], self.valid_padding[1]), \
+ (self.valid_padding[2], self.valid_padding[2])), mode='reflect')
+ self.labels[k] = np.pad(self.labels[k], ((self.valid_padding[0], self.valid_padding[0]), \
+ (self.valid_padding[1], self.valid_padding[1]), \
+ (self.valid_padding[2], self.valid_padding[2])), mode='reflect')
+
+ # the training dataset size
+ self.raw_data_shape = list(self.dataset[0].shape)
+ print('valid_size:', self.raw_data_shape)
+
+ self.reset_output()
+ self.weight_vol = self.get_weight()
+
+ # the number of inference times
+ self.num_per_dataset = self.num_zyx[0] * self.num_zyx[1] * self.num_zyx[2]
+ self.iters_num = self.num_per_dataset * len(self.dataset)
+
+ def __getitem__(self, index):
+ # print(index)
+ pos_data = index // self.num_per_dataset
+ pre_data = index % self.num_per_dataset
+ pos_z = pre_data // (self.num_zyx[1] * self.num_zyx[2])
+ pos_xy = pre_data % (self.num_zyx[1] * self.num_zyx[2])
+ pos_x = pos_xy // self.num_zyx[2]
+ pos_y = pos_xy % self.num_zyx[2]
+
+ # find position
+ fromz = pos_z * self.stride[0]
+ endz = fromz + self.crop_size[0]
+ if endz > self.raw_data_shape[0]:
+ endz = self.raw_data_shape[0]
+ fromz = endz - self.crop_size[0]
+ fromy = pos_y * self.stride[1]
+ endy = fromy + self.crop_size[1]
+ if endy > self.raw_data_shape[1]:
+ endy = self.raw_data_shape[1]
+ fromy = endy - self.crop_size[1]
+ fromx = pos_x * self.stride[2]
+ endx = fromx + self.crop_size[2]
+ if endx > self.raw_data_shape[2]:
+ endx = self.raw_data_shape[2]
+ fromx = endx - self.crop_size[2]
+
+ self.pos = [fromz, fromy, fromx]
+
+ imgs = self.dataset[pos_data][fromz:endz, fromx:endx, fromy:endy].copy()
+ lb = self.labels[pos_data][fromz:endz, fromx:endx, fromy:endy].copy()
+ lb_affs = seg_to_affgraph(lb, mknhood3d(1), pad='replicate').astype(np.float32)
+
+ weight_factor = np.sum(lb_affs) / np.size(lb_affs)
+ weight_factor = np.clip(weight_factor, 1e-3, 1)
+ weightmap = lb_affs * (1 - weight_factor) / weight_factor + (1 - lb_affs)
+
+ imgs = imgs.astype(np.float32) / 255.0
+ imgs = imgs[np.newaxis, ...]
+ imgs = np.ascontiguousarray(imgs, dtype=np.float32)
+ lb_affs = np.ascontiguousarray(lb_affs, dtype=np.float32)
+ weightmap = np.ascontiguousarray(weightmap, dtype=np.float32)
+ return imgs, lb_affs, weightmap
+
+ def __len__(self):
+ return self.iters_num
+
+ def reset_output(self):
+ self.out_affs = np.zeros(tuple([3] + self.raw_data_shape), dtype=np.float32)
+ self.weight_map = np.zeros(tuple([1] + self.raw_data_shape), dtype=np.float32)
+
+ self.out_affs2 = np.zeros(tuple([1] + self.raw_data_shape), dtype=np.float32)
+ self.weight_map2 = np.zeros(tuple([1] + self.raw_data_shape), dtype=np.float32)
+
+ self.out_affs3 = np.zeros(tuple([1] + self.raw_data_shape), dtype=np.float32)
+ self.weight_map3 = np.zeros(tuple([1] + self.raw_data_shape), dtype=np.float32)
+
+ def get_weight(self, sigma=0.2, mu=0.0):
+ zz, yy, xx = np.meshgrid(np.linspace(-1, 1, self.out_size[0], dtype=np.float32),
+ np.linspace(-1, 1, self.out_size[1], dtype=np.float32),
+ np.linspace(-1, 1, self.out_size[2], dtype=np.float32), indexing='ij')
+ dd = np.sqrt(zz * zz + yy * yy + xx * xx)
+ weight = 1e-6 + np.exp(-((dd - mu) ** 2 / (2.0 * sigma ** 2)))
+ weight = weight[np.newaxis, ...]
+ return weight
+
+ def add_vol(self, affs_vol):
+ fromz, fromy, fromx = self.pos
+ self.out_affs[:, fromz:fromz + self.out_size[0], \
+ fromx:fromx + self.out_size[1], \
+ fromy:fromy + self.out_size[2]] += affs_vol * self.weight_vol
+ self.weight_map[:, fromz:fromz + self.out_size[0], \
+ fromx:fromx + self.out_size[1], \
+ fromy:fromy + self.out_size[2]] += self.weight_vol
+
+ def get_results(self):
+ self.out_affs = self.out_affs / self.weight_map
+ self.out_affs = self.out_affs[:, self.valid_padding[0]:-self.valid_padding[0], \
+ self.valid_padding[1]:-self.valid_padding[1], \
+ self.valid_padding[2]:-self.valid_padding[2]]
+
+ return self.out_affs
+
+ def add_bound(self, affs_vol):
+ fromz, fromy, fromx = self.pos
+ self.out_affs2[:, fromz:fromz + self.out_size[0], \
+ fromx:fromx + self.out_size[1], \
+ fromy:fromy + self.out_size[2]] += affs_vol * self.weight_vol
+ self.weight_map2[:, fromz:fromz + self.out_size[0], \
+ fromx:fromx + self.out_size[1], \
+ fromy:fromy + self.out_size[2]] += self.weight_vol
+
+ def get_results_bound(self):
+ self.out_affs2 = self.out_affs2 / self.weight_map2
+ self.out_affs2 = self.out_affs2[:, self.valid_padding[0]:-self.valid_padding[0], \
+ self.valid_padding[1]:-self.valid_padding[1], \
+ self.valid_padding[2]:-self.valid_padding[2]]
+
+ return self.out_affs2
+
+ def get_gt_affs(self, num_data=0):
+ return self.gt_affs[num_data].copy()
+
+ def get_gt_lb(self, num_data=0):
+ lbs = self.labels[num_data].copy()
+ return lbs[self.valid_padding[0]:-self.valid_padding[0], \
+ self.valid_padding[1]:-self.valid_padding[1], \
+ self.valid_padding[2]:-self.valid_padding[2]]
+
+ def get_raw_data(self, num_data=0):
+ out = self.dataset[num_data].copy()
+ return out[self.valid_padding[0]:-self.valid_padding[0], \
+ self.valid_padding[1]:-self.valid_padding[1], \
+ self.valid_padding[2]:-self.valid_padding[2]]
diff --git a/legacy/Train_and_Inference/loss/loss.py b/legacy/Train_and_Inference/loss/loss.py
new file mode 100644
index 0000000..21a49b3
--- /dev/null
+++ b/legacy/Train_and_Inference/loss/loss.py
@@ -0,0 +1,190 @@
+from __future__ import print_function, division
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+#######################################################
+# 0. Main loss functions
+#######################################################
+
+class JaccardLoss(nn.Module):
+ """Jaccard loss.
+ """
+ # binary case
+
+ def __init__(self, size_average=True, reduce=True, smooth=1.0):
+ super(JaccardLoss, self).__init__()
+ self.smooth = smooth
+ self.reduce = reduce
+
+ def jaccard_loss(self, pred, target):
+ loss = 0.
+ # for each sample in the batch
+ for index in range(pred.size()[0]):
+ iflat = pred[index].view(-1)
+ tflat = target[index].view(-1)
+ intersection = (iflat * tflat).sum()
+ loss += 1 - ((intersection + self.smooth) /
+ ( iflat.sum() + tflat.sum() - intersection + self.smooth))
+ #print('loss:',intersection, iflat.sum(), tflat.sum())
+
+ # size_average=True for the jaccard loss
+ return loss / float(pred.size()[0])
+
+ def jaccard_loss_batch(self, pred, target):
+ iflat = pred.view(-1)
+ tflat = target.view(-1)
+ intersection = (iflat * tflat).sum()
+ loss = 1 - ((intersection + self.smooth) /
+ ( iflat.sum() + tflat.sum() - intersection + self.smooth))
+ #print('loss:',intersection, iflat.sum(), tflat.sum())
+ return loss
+
+ def forward(self, pred, target):
+ #_assert_no_grad(target)
+ if not (target.size() == pred.size()):
+ raise ValueError("Target size ({}) must be the same as pred size ({})".format(target.size(), pred.size()))
+ if self.reduce:
+ loss = self.jaccard_loss(pred, target)
+ else:
+ loss = self.jaccard_loss_batch(pred, target)
+ return loss
+
+class DiceLoss(nn.Module):
+ """DICE loss.
+ """
+ # https://lars76.github.io/neural-networks/object-detection/losses-for-segmentation/
+
+ def __init__(self, size_average=True, reduce=True, smooth=100.0, power=1):
+ super(DiceLoss, self).__init__()
+ self.smooth = smooth
+ self.reduce = reduce
+ self.power = power
+
+ def dice_loss(self, pred, target):
+ loss = 0.
+
+ for index in range(pred.size()[0]):
+ iflat = pred[index].view(-1)
+ tflat = target[index].view(-1)
+ intersection = (iflat * tflat).sum()
+ if self.power==1:
+ loss += 1 - ((2. * intersection + self.smooth) /
+ ( iflat.sum() + tflat.sum() + self.smooth))
+ else:
+ loss += 1 - ((2. * intersection + self.smooth) /
+ ( (iflat**self.power).sum() + (tflat**self.power).sum() + self.smooth))
+
+ # size_average=True for the dice loss
+ return loss / float(pred.size()[0])
+
+ def dice_loss_batch(self, pred, target):
+ iflat = pred.view(-1)
+ tflat = target.view(-1)
+ intersection = (iflat * tflat).sum()
+
+ if self.power==1:
+ loss = 1 - ((2. * intersection + self.smooth) /
+ (iflat.sum() + tflat.sum() + self.smooth))
+ else:
+ loss = 1 - ((2. * intersection + self.smooth) /
+ ( (iflat**self.power).sum() + (tflat**self.power).sum() + self.smooth))
+ return loss
+
+ def forward(self, pred, target):
+ #_assert_no_grad(target)
+ if not (target.size() == pred.size()):
+ raise ValueError("Target size ({}) must be the same as pred size ({})".format(target.size(), pred.size()))
+
+ if self.reduce:
+ loss = self.dice_loss(pred, target)
+ else:
+ loss = self.dice_loss_batch(pred, target)
+ return loss
+
+class WeightedMSE(nn.Module):
+ """Weighted mean-squared error.
+ """
+
+ def __init__(self):
+ super().__init__()
+
+ def weighted_mse_loss(self, pred, target, weight):
+
+ s1 = torch.prod(torch.tensor(pred.size()[2:]).float())
+ s2 = pred.size()[0]
+ norm_term = (s1 * s2).cuda()
+ if weight is None:
+ return torch.sum((pred - target) ** 2) / norm_term
+ else:
+ return torch.sum(weight * (pred - target) ** 2) / norm_term
+
+ def forward(self, pred, target, weight=None):
+ #_assert_no_grad(target)
+ return self.weighted_mse_loss(pred, target, weight)
+
+class MSELoss(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.criterion = nn.MSELoss()
+
+ def forward(self, pred, target, weight=None):
+ return self.criterion(pred, target)
+
+class BCELoss(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.criterion = nn.BCELoss()
+
+ def forward(self, pred, target, weight=None):
+ return self.criterion(pred, target)
+
+class WeightedBCE(nn.Module):
+ """Weighted binary cross-entropy.
+ """
+ def __init__(self, size_average=True, reduce=True):
+ super().__init__()
+ self.size_average = size_average
+ self.reduce = reduce
+
+ def forward(self, pred, target, weight=None):
+ #_assert_no_grad(target)
+ if pred.min()<0:
+ # pred = (pred-pred.min())/(pred-pred.min()).max()
+ pred = F.relu(pred)
+
+ return F.binary_cross_entropy(pred, target, weight)
+
+class WeightedCE(nn.Module):
+ """Mask weighted multi-class cross-entropy (CE) loss.
+ """
+ def __init__(self):
+ super().__init__()
+
+ def forward(self, pred, target, weight_mask=None):
+ # Different from, F.binary_cross_entropy, the "weight" parameter
+ # in F.cross_entropy is a manual rescaling weight given to each
+ # class. Therefore we need to multiply the weight mask after the
+ # loss calculation.
+ loss = F.cross_entropy(pred, target, reduction='none')
+ if weight_mask is not None:
+ loss = loss * weight_mask
+ return loss.mean()
+
+#######################################################
+# 1. Regularization
+#######################################################
+
+class BinaryReg(nn.Module):
+ """Regularization for encouraging the outputs to be binary.
+ """
+ def __init__(self, alpha=0.1):
+ super().__init__()
+ self.alpha = alpha
+
+ def forward(self, pred):
+ diff = pred - 0.5
+ diff = torch.clamp(torch.abs(diff), min=1e-2)
+ loss = (1.0 / diff).mean()
+ return self.alpha * loss
diff --git a/legacy/Train_and_Inference/model/Mnet.py b/legacy/Train_and_Inference/model/Mnet.py
new file mode 100644
index 0000000..3c535eb
--- /dev/null
+++ b/legacy/Train_and_Inference/model/Mnet.py
@@ -0,0 +1,310 @@
+from torch import nn
+from torch import nn
+import torch
+import torch.nn.functional as F
+
+
+class CNA3d(nn.Module): # conv + norm + activation
+ def __init__(self, in_channels, out_channels, kSize, stride, padding=(1, 1, 1), bias=True, norm_args=None,
+ activation_args=None):
+ super().__init__()
+ self.norm_args = norm_args
+ self.activation_args = activation_args
+
+ self.conv = nn.Conv3d(in_channels, out_channels, kernel_size=kSize, stride=stride, padding=padding, bias=bias)
+
+ if norm_args is not None:
+ self.norm = nn.InstanceNorm3d(out_channels, **norm_args)
+
+ if activation_args is not None:
+ self.activation = nn.LeakyReLU(**activation_args)
+
+ def forward(self, x):
+ x = self.conv(x)
+
+ if self.norm_args is not None:
+ x = self.norm(x)
+
+ if self.activation_args is not None:
+ x = self.activation(x)
+ return x
+
+
+class CB3d(nn.Module): # conv block 3d
+ def __init__(self, in_channels, out_channels, kSize=(3, 3), stride=(1, 1), padding=(1, 1, 1), bias=True,
+ norm_args: tuple = (None, None), activation_args: tuple = (None, None)):
+ super().__init__()
+
+ self.conv1 = CNA3d(in_channels, out_channels, kSize=kSize[0], stride=stride[0],
+ padding=padding, bias=bias, norm_args=norm_args[0], activation_args=activation_args[0])
+
+ self.conv2 = CNA3d(out_channels, out_channels, kSize=kSize[1], stride=stride[1],
+ padding=padding, bias=bias, norm_args=norm_args[1], activation_args=activation_args[1])
+
+ def forward(self, x):
+ x = self.conv1(x)
+ x = self.conv2(x)
+ return x
+
+
+class BasicNet(nn.Module):
+ norm_kwargs = {'affine': True}
+ activation_kwargs = {'negative_slope': 1e-2, 'inplace': True}
+
+ def __init__(self):
+ super(BasicNet, self).__init__()
+
+ def parameter_count(self):
+ print("model have {} paramerters in total".format(sum(x.numel() for x in self.parameters()) / 1e6))
+
+
+def FMU(x1, x2, mode='sub'):
+ """
+ feature merging unit
+ Args:
+ x1:
+ x2:
+ mode: type of fusion
+ Returns:
+ """
+ if mode == 'sum':
+ return torch.add(x1, x2)
+ elif mode == 'sub':
+ return torch.abs(x1 - x2)
+ elif mode == 'cat':
+ return torch.cat((x1, x2), dim=1)
+ else:
+ raise Exception('Unexpected mode')
+
+
+class Down(BasicNet):
+ def __init__(self, in_channels, out_channels, mode: tuple, FMU='sub', downsample=True, min_z=8):
+ """
+ basic module at downsampling stage
+ Args:
+ in_channels:
+ out_channels:
+ mode: represent the streams coming in and out. e.g., ('2d', 'both'): one input stream (2d) and two output streams (2d and 3d)
+ FMU: determine the type of feature fusion if there are two input streams
+ downsample: determine whether to downsample input features (only the first module of MNet do not downsample)
+ min_z: if the size of z-axis < min_z, maxpooling won't be applied along z-axis
+ """
+ super().__init__()
+ self.mode_in, self.mode_out = mode
+ self.downsample = downsample
+ self.FMU = FMU
+ self.min_z = min_z
+ norm_args = (self.norm_kwargs, self.norm_kwargs)
+ activation_args = (self.activation_kwargs, self.activation_kwargs)
+
+ if self.mode_out == '2d' or self.mode_out == 'both':
+ self.CB2d = CB3d(in_channels=in_channels, out_channels=out_channels,
+ kSize=((1, 3, 3), (1, 3, 3)), stride=(1, 1), padding=(0, 1, 1),
+ norm_args=norm_args, activation_args=activation_args)
+
+ if self.mode_out == '3d' or self.mode_out == 'both':
+ self.CB3d = CB3d(in_channels=in_channels, out_channels=out_channels,
+ kSize=(3, 3), stride=(1, 1), padding=(1, 1, 1),
+ norm_args=norm_args, activation_args=activation_args)
+
+ def forward(self, x):
+ if self.downsample:
+ if self.mode_in == 'both':
+ x2d, x3d = x
+ p2d = F.max_pool3d(x2d, kernel_size=(1, 2, 2), stride=(1, 2, 2))
+ if x3d.shape[2] >= self.min_z:
+ p3d = F.max_pool3d(x3d, kernel_size=(2, 2, 2), stride=(2, 2, 2))
+ else:
+ p3d = F.max_pool3d(x3d, kernel_size=(1, 2, 2), stride=(1, 2, 2))
+
+ x = FMU(p2d, p3d, mode=self.FMU)
+
+ elif self.mode_in == '2d':
+ x = F.max_pool3d(x, kernel_size=(1, 2, 2), stride=(1, 2, 2))
+
+ elif self.mode_in == '3d':
+ if x.shape[2] >= self.min_z:
+ x = F.max_pool3d(x, kernel_size=(2, 2, 2), stride=(2, 2, 2))
+ else:
+ x = F.max_pool3d(x, kernel_size=(1, 2, 2), stride=(1, 2, 2))
+
+ if self.mode_out == '2d':
+ return self.CB2d(x)
+ elif self.mode_out == '3d':
+ return self.CB3d(x)
+ elif self.mode_out == 'both':
+ return self.CB2d(x), self.CB3d(x)
+
+
+class Up(BasicNet):
+ def __init__(self, in_channels, out_channels, mode: tuple, FMU='sub'):
+ """
+ basic module at upsampling stage
+ Args:
+ in_channels:
+ out_channels:
+ mode: represent the streams coming in and out. e.g., ('2d', 'both'): one input stream (2d) and two output streams (2d and 3d)
+ FMU: determine the type of feature fusion if there are two input streams
+ """
+ super().__init__()
+ self.mode_in, self.mode_out = mode
+ self.FMU = FMU
+ norm_args = (self.norm_kwargs, self.norm_kwargs)
+ activation_args = (self.activation_kwargs, self.activation_kwargs)
+
+ if self.mode_out == '2d' or self.mode_out == 'both':
+ self.CB2d = CB3d(in_channels=in_channels, out_channels=out_channels,
+ kSize=((1, 3, 3), (1, 3, 3)), stride=(1, 1), padding=(0, 1, 1),
+ norm_args=norm_args, activation_args=activation_args)
+
+ if self.mode_out == '3d' or self.mode_out == 'both':
+ self.CB3d = CB3d(in_channels=in_channels, out_channels=out_channels,
+ kSize=(3, 3), stride=(1, 1), padding=(1, 1, 1),
+ norm_args=norm_args, activation_args=activation_args)
+
+ def forward(self, x):
+ x2d, xskip2d, x3d, xskip3d = x
+
+ tarSize = xskip2d.shape[2:]
+ up2d = F.interpolate(x2d, size=tarSize, mode='trilinear', align_corners=False)
+ up3d = F.interpolate(x3d, size=tarSize, mode='trilinear', align_corners=False)
+
+ cat = torch.cat([FMU(xskip2d, xskip3d, self.FMU), FMU(up2d, up3d, self.FMU)], dim=1)
+
+ if self.mode_out == '2d':
+ return self.CB2d(cat)
+ elif self.mode_out == '3d':
+ return self.CB3d(cat)
+ elif self.mode_out == 'both':
+ return self.CB2d(cat), self.CB3d(cat)
+
+
+class MNet(BasicNet):
+ def __init__(self, in_channels, kn=(32, 48, 64, 80, 96), FMU='sub'):
+ """
+
+ Args:
+ in_channels: channels of input
+ num_classes: output classes
+ kn: the number of kernels
+ ds: deep supervision
+ FMU: type of feature merging unit
+ """
+ super().__init__()
+
+ channel_factor = {'sum': 1, 'sub': 1, 'cat': 2}
+ fct = channel_factor[FMU]
+
+ self.down11 = Down(in_channels, kn[0], ('/', 'both'), downsample=False)
+ self.down12 = Down(kn[0], kn[1], ('2d', 'both'))
+ self.down13 = Down(kn[1], kn[2], ('2d', 'both'))
+ self.down14 = Down(kn[2], kn[3], ('2d', 'both'))
+ self.bottleneck1 = Down(kn[3], kn[4], ('2d', '2d'))
+ self.up11 = Up(fct * (kn[3] + kn[4]), kn[3], ('both', '2d'), FMU)
+ self.up12 = Up(fct * (kn[2] + kn[3]), kn[2], ('both', '2d'), FMU)
+ self.up13 = Up(fct * (kn[1] + kn[2]), kn[1], ('both', '2d'), FMU)
+ self.up14 = Up(fct * (kn[0] + kn[1]), kn[0], ('both', 'both'), FMU)
+
+ self.up11_ = Up(fct * (kn[3] + kn[4]), kn[3], ('both', '2d'), FMU)
+ self.up12_ = Up(fct * (kn[2] + kn[3]), kn[2], ('both', '2d'), FMU)
+ self.up13_ = Up(fct * (kn[1] + kn[2]), kn[1], ('both', '2d'), FMU)
+ self.up14_ = Up(fct * (kn[0] + kn[1]), kn[0], ('both', 'both'), FMU)
+
+ self.down21 = Down(kn[0], kn[1], ('3d', 'both'))
+ self.down22 = Down(fct * kn[1], kn[2], ('both', 'both'), FMU)
+ self.down23 = Down(fct * kn[2], kn[3], ('both', 'both'), FMU)
+ self.bottleneck2 = Down(fct * kn[3], kn[4], ('both', 'both'), FMU)
+ self.up21 = Up(fct * (kn[3] + kn[4]), kn[3], ('both', 'both'), FMU)
+ self.up22 = Up(fct * (kn[2] + kn[3]), kn[2], ('both', 'both'), FMU)
+ self.up23 = Up(fct * (kn[1] + kn[2]), kn[1], ('both', '3d'), FMU)
+
+ self.up21_ = Up(fct * (kn[3] + kn[4]), kn[3], ('both', 'both'), FMU)
+ self.up22_ = Up(fct * (kn[2] + kn[3]), kn[2], ('both', 'both'), FMU)
+ self.up23_ = Up(fct * (kn[1] + kn[2]), kn[1], ('both', '3d'), FMU)
+
+ self.down31 = Down(kn[1], kn[2], ('3d', 'both'))
+ self.down32 = Down(fct * kn[2], kn[3], ('both', 'both'), FMU)
+ self.bottleneck3 = Down(fct * kn[3], kn[4], ('both', 'both'), FMU)
+ self.up31 = Up(fct * (kn[3] + kn[4]), kn[3], ('both', 'both'), FMU)
+ self.up32 = Up(fct * (kn[2] + kn[3]), kn[2], ('both', '3d'), FMU)
+
+ self.up31_ = Up(fct * (kn[3] + kn[4]), kn[3], ('both', 'both'), FMU)
+ self.up32_ = Up(fct * (kn[2] + kn[3]), kn[2], ('both', '3d'), FMU)
+
+ self.down41 = Down(kn[2], kn[3], ('3d', 'both'), FMU)
+ self.bottleneck4 = Down(fct * kn[3], kn[4], ('both', 'both'), FMU)
+ self.up41 = Up(fct * (kn[3] + kn[4]), kn[3], ('both', '3d'), FMU)
+
+ self.up41_ = Up(fct * (kn[3] + kn[4]), kn[3], ('both', '3d'), FMU)
+
+ self.bottleneck5 = Down(kn[3], kn[4], ('3d', '3d'))
+
+ self.outputs = nn.ModuleList(
+ [nn.Conv3d(c, 3, kernel_size=(1, 1, 1), stride=1, padding=0, bias=False)
+ for c in [kn[0], kn[1], kn[1], kn[2], kn[2], kn[3], kn[3]]]
+ )
+
+ self.outputs2 = nn.ModuleList(
+ [nn.Conv3d(c, 1, kernel_size=(1, 1, 1), stride=1, padding=0, bias=False)
+ for c in [kn[0], kn[1], kn[1], kn[2], kn[2], kn[3], kn[3]]]
+ )
+ self.sigmoid = nn.Sigmoid()
+
+ def forward(self, x):
+ down11 = self.down11(x)
+ down12 = self.down12(down11[0])
+ down13 = self.down13(down12[0])
+ down14 = self.down14(down13[0])
+ bottleNeck1 = self.bottleneck1(down14[0])
+
+ down21 = self.down21(down11[1])
+ down22 = self.down22([down21[0], down12[1]])
+ down23 = self.down23([down22[0], down13[1]])
+ bottleNeck2 = self.bottleneck2([down23[0], down14[1]])
+
+ down31 = self.down31(down21[1])
+ down32 = self.down32([down31[0], down22[1]])
+ bottleNeck3 = self.bottleneck3([down32[0], down23[1]])
+
+ down41 = self.down41(down31[1])
+ bottleNeck4 = self.bottleneck4([down41[0], down32[1]])
+
+ bottleNeck5 = self.bottleneck5(down41[1])
+
+ up41 = self.up41([bottleNeck4[0], down41[0], bottleNeck5, down41[1]])
+
+ up31 = self.up31([bottleNeck3[0], down32[0], bottleNeck4[1], down32[1]])
+ up32 = self.up32([up31[0], down31[0], up41, down31[1]])
+
+ up21 = self.up21([bottleNeck2[0], down23[0], bottleNeck3[1], down23[1]])
+ up22 = self.up22([up21[0], down22[0], up31[1], down22[1]])
+ up23 = self.up23([up22[0], down21[0], up32, down21[1]])
+
+ up11 = self.up11([bottleNeck1, down14[0], bottleNeck2[1], down14[1]])
+ up12 = self.up12([up11, down13[0], up21[1], down13[1]])
+ up13 = self.up13([up12, down12[0], up22[1], down12[1]])
+ up14 = self.up14([up13, down11[0], up23, down11[1]])
+
+ up41_ = self.up41_([bottleNeck4[0], down41[0], bottleNeck5, down41[1]])
+
+ up31_ = self.up31_([bottleNeck3[0], down32[0], bottleNeck4[1], down32[1]])
+ up32_ = self.up32_([up31_[0], down31[0], up41_, down31[1]])
+
+ up21_ = self.up21_([bottleNeck2[0], down23[0], bottleNeck3[1], down23[1]])
+ up22_ = self.up22_([up21_[0], down22[0], up31_[1], down22[1]])
+ up23_ = self.up23_([up22_[0], down21[0], up32_, down21[1]])
+
+ up11_ = self.up11_([bottleNeck1, down14[0], bottleNeck2[1], down14[1]])
+ up12_ = self.up12_([up11_, down13[0], up21_[1], down13[1]])
+ up13_ = self.up13_([up12_, down12[0], up22_[1], down12[1]])
+ up14_ = self.up14_([up13_, down11[0], up23_, down11[1]])
+
+ return self.sigmoid(self.outputs[0](up14[0] + up14[1])), self.sigmoid(self.outputs2[0](up14_[0] + up14_[1]))
+
+
+if __name__ == '__main__':
+ MNet = MNet(1, kn=(28, 36, 48, 64, 80), FMU='sub')
+ input = torch.randn((1, 1, 32, 96, 96))
+ output = MNet(input)
+
+ print([e.shape for e in output])
diff --git a/legacy/Train_and_Inference/supervised_provider.py b/legacy/Train_and_Inference/supervised_provider.py
new file mode 100644
index 0000000..29884d8
--- /dev/null
+++ b/legacy/Train_and_Inference/supervised_provider.py
@@ -0,0 +1,380 @@
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+
+import os
+import sys
+import random
+import numpy as np
+from torch.utils.data import Dataset
+from torch.utils.data import DataLoader
+from utils.augmentation import SimpleAugment as Filp
+from utils.consistency_aug_perturbations import Intensity
+from utils.consistency_aug_perturbations import GaussBlur
+from utils.consistency_aug_perturbations import GaussNoise
+from utils.consistency_aug_perturbations_sup import Cutout
+from utils.augmentation import ElasticAugment as Elastic
+import imageio
+from utils.seg_util import mknhood3d
+from utils.aff_util import seg_to_affgraph
+
+
+class Train(Dataset):
+ def __init__(self, cfg):
+ super(Train, self).__init__()
+ self.cfg = cfg
+ self.model_type = cfg.MODEL.model_type
+
+ # split training data
+ self.start_slice = cfg.DATA.start_slice
+ self.end_slice = cfg.DATA.end_slice
+ # augmentation
+ self.simple_aug = Filp()
+ self.min_noise_std = cfg.DATA.min_noise_std
+ self.max_noise_std = cfg.DATA.max_noise_std
+ self.min_kernel_size = cfg.DATA.min_kernel_size
+ self.max_kernel_size = cfg.DATA.max_kernel_size
+ self.min_sigma = cfg.DATA.min_sigma
+ self.max_sigma = cfg.DATA.max_sigma
+ self.perturbations_init()
+ self.crop_from_origin = [20, 128, 128]
+
+ self.dataset = []
+ self.labels = []
+ dataset_list = ['J0126-sbem', 'Kasthuri-atum', 'Hemi-brain-fib', 'CREMI-sstem', 'AxonEM[M]-sstem',
+ 'AxonEM[H]-atum', 'Mira-adwt', 'Fib-25-fib', 'Mira-scn', 'Mira-fish']
+
+ for sub_path in dataset_list:
+ self.folder_name = os.path.join(cfg.DATA.data_folder, sub_path)
+ file_num = len(os.listdir(self.folder_name)) // 2
+ train_datasets = ['%d.tif' % i for i in range(file_num)]
+ train_labels = ['%d_MaskIns.tif' % i for i in range(file_num)]
+
+ for k in range(len(train_datasets)):
+ print('load ' + self.folder_name + train_datasets[k] + ' ...')
+ data = imageio.volread(os.path.join(self.folder_name, train_datasets[k]))
+ self.dataset.append(data[:])
+ label = imageio.volread(os.path.join(self.folder_name, train_labels[k]))
+ self.labels.append(label[:])
+
+ print(len(self.dataset))
+
+ def __getitem__(self, index):
+
+ # dataset-balanced sampling
+ type_dataset = random.randint(0, 9)
+ if type_dataset == 0:
+ k = random.randint(0, 32) # j0126
+ elif type_dataset == 1:
+ k = random.randint(33, 34) # kstr
+ elif type_dataset == 2:
+ k = random.randint(35, 40) # hemibrain
+ elif type_dataset == 3:
+ k = random.randint(41, 43) # cremi
+ elif type_dataset == 4:
+ k = random.randint(44, 52) # axon
+ elif type_dataset == 5:
+ k = random.randint(53, 61) # axon
+ elif type_dataset == 6:
+ k = random.randint(62, 64) # adwt
+ elif type_dataset == 7:
+ k = random.randint(65, 66) # fib25
+ elif type_dataset == 8:
+ k = random.randint(67, 71) # scn
+ elif type_dataset == 9:
+ k = random.randint(72, 75) # fish
+
+ # multi-z-ratio aug.
+ if type_dataset == 2 or type_dataset == 7:
+ z_scale_ratio = random.randint(1, 3)
+ used_data = self.dataset[k][::z_scale_ratio, :, :]
+ used_label = self.labels[k][::z_scale_ratio, :, :]
+ elif type_dataset == 0:
+ z_scale_ratio = random.randint(1, 2)
+ used_data = self.dataset[k][::z_scale_ratio, :, :]
+ used_label = self.labels[k][::z_scale_ratio, :, :]
+ else:
+ used_data = self.dataset[k]
+ used_label = self.labels[k]
+
+ raw_data_shape = used_data.shape
+
+ random_z = random.randint(0, raw_data_shape[0] - self.crop_from_origin[0])
+ random_y = random.randint(0, raw_data_shape[1] - self.crop_from_origin[1])
+ random_x = random.randint(0, raw_data_shape[2] - self.crop_from_origin[2])
+
+ imgs1 = used_data[random_z:random_z + self.crop_from_origin[0], \
+ random_y:random_y + self.crop_from_origin[1], \
+ random_x:random_x + self.crop_from_origin[2]].copy()
+ lb1 = used_label[random_z:random_z + self.crop_from_origin[0], \
+ random_y:random_y + self.crop_from_origin[1], \
+ random_x:random_x + self.crop_from_origin[2]].copy()
+
+ imgs1 = imgs1.astype(np.float32) / 255.0
+
+ # style-mixing
+ if random.random() < self.cfg.TRAIN.freq_mix_prob:
+ type_dataset_r = random.choice(list(range(0, type_dataset)) + list(range(type_dataset + 1, 9)))
+ if type_dataset_r == 0:
+ r = random.randint(0, 32) # j0126
+ elif type_dataset_r == 1:
+ r = random.randint(33, 34) # kstr
+ elif type_dataset_r == 2:
+ r = random.randint(35, 40) # hemibrain
+ elif type_dataset_r == 3:
+ r = random.randint(41, 43) # cremi
+ elif type_dataset_r == 4:
+ r = random.randint(44, 52) # axon
+ elif type_dataset_r == 5:
+ r = random.randint(53, 61) # axon
+ elif type_dataset_r == 6:
+ r = random.randint(62, 64) # adwt
+ elif type_dataset_r == 7:
+ r = random.randint(65, 66) # fib25
+ elif type_dataset_r == 8:
+ r = random.randint(67, 71) # scn
+ elif type_dataset_r == 9:
+ r = random.randint(72, 75) # fish
+
+ used_data_r = self.dataset[r]
+ raw_data_shape_r = used_data_r.shape
+ random_z = random.randint(0, raw_data_shape_r[0] - self.crop_from_origin[0])
+ random_y = random.randint(0, raw_data_shape_r[1] - self.crop_from_origin[1])
+ random_x = random.randint(0, raw_data_shape_r[2] - self.crop_from_origin[2])
+
+ imgs2 = used_data_r[random_z:random_z + self.crop_from_origin[0], \
+ random_y:random_y + self.crop_from_origin[1], \
+ random_x:random_x + self.crop_from_origin[2]].copy()
+ imgs2 = imgs2.astype(np.float32) / 255.0
+ imgs1 = FDA_source_to_target_np_3D(imgs1, imgs2, 0.001)
+
+ imgs1[imgs1 < 0] = 0
+ imgs1[imgs1 > 1] = 1
+
+ # normal augumentation
+ [imgs1, lb1] = self.simple_aug([imgs1, lb1])
+ imgs1, lb1 = self.apply_perturbations(imgs1, lb1)
+
+ lb_affs1 = seg_to_affgraph(lb1, mknhood3d(1), pad='replicate').astype(np.float32)
+ bounds1 = np.float32(lb1 != 0)
+
+ type_dataset_r2, r2 = -1, -1
+ if random.random() < self.cfg.TRAIN.spa_mix_prob:
+ if random.random() < 0.5:
+ type_dataset_r2 = random.choice((list(range(0, type_dataset)) + list(range(type_dataset + 1, 9))))
+ if type_dataset_r2 == 0:
+ r2 = random.randint(0, 32) # j0126
+ elif type_dataset_r2 == 1:
+ r2 = random.randint(33, 34) # kstr
+ elif type_dataset_r2 == 2:
+ r2 = random.randint(35, 40) # hemibrain
+ elif type_dataset_r2 == 3:
+ r2 = random.randint(41, 43) # cremi
+ elif type_dataset_r2 == 4:
+ r2 = random.randint(44, 52) # axon
+ elif type_dataset_r2 == 5:
+ r2 = random.randint(53, 61) # axon
+ elif type_dataset_r2 == 6:
+ r2 = random.randint(62, 64) # adwt
+ elif type_dataset_r2 == 7:
+ r2 = random.randint(65, 66) # fib25
+ elif type_dataset_r2 == 8:
+ r2 = random.randint(67, 71) # scn
+ elif type_dataset_r2 == 9:
+ r2 = random.randint(72, 75) # fish
+
+ used_data = self.dataset[r2]
+ used_label = self.labels[r2]
+ raw_data_shape = used_data.shape
+
+ random_z = random.randint(0, raw_data_shape[0] - self.crop_from_origin[0])
+ random_y = random.randint(0, raw_data_shape[1] - self.crop_from_origin[1])
+ random_x = random.randint(0, raw_data_shape[2] - self.crop_from_origin[2])
+ imgs2 = used_data[random_z:random_z + self.crop_from_origin[0], \
+ random_y:random_y + self.crop_from_origin[1], \
+ random_x:random_x + self.crop_from_origin[2]].copy()
+ lb2 = used_label[random_z:random_z + self.crop_from_origin[0], \
+ random_y:random_y + self.crop_from_origin[1], \
+ random_x:random_x + self.crop_from_origin[2]].copy()
+ imgs2 = imgs2.astype(np.float32) / 255.0
+ [imgs2, lb2] = self.simple_aug([imgs2, lb2])
+ imgs2, lb2 = self.apply_perturbations(imgs2, lb2)
+
+ lb_affs2 = seg_to_affgraph(lb2, mknhood3d(1), pad='replicate').astype(np.float32)
+ bounds2 = np.float32(lb2 != 0)
+
+ random_y_mask = random.randint(0, 58)
+ random_x_mask = random.randint(0, 58)
+ mask = np.ones(self.crop_from_origin)
+ mask[:, random_y_mask:random_y_mask + 70, random_x_mask:random_x_mask + 70] = 0
+
+ imgs = mask * imgs1 + (1 - mask) * imgs2
+ lb_affs = np.array([mask, mask, mask]) * lb_affs1 + (1 - np.array([mask, mask, mask])) * lb_affs2
+ lb = mask * lb1 + (1 - mask) * lb2
+ bounds = mask * bounds1 + (1 - mask) * bounds2
+
+ else:
+ imgs = imgs1
+ lb_affs = lb_affs1
+ lb = lb1
+ bounds = bounds1
+
+ # area mask
+ if True:
+ if type_dataset == 10 or type_dataset == 0 or type_dataset == 8 or type_dataset_r2 == 0 or type_dataset_r2 == 8:
+ bmask = 1 - np.uint(lb == 10000)
+ elif (type_dataset == 1 and k == 33) or (type_dataset_r2 == 1 and r2 == 33):
+ bmask = 1 - (np.uint(lb == 121) + np.uint(lb == 122))
+ else:
+ bmask = np.ones_like(lb)
+ else:
+ bmask = np.ones_like(lb)
+
+ # extend dimension
+ imgs = imgs[np.newaxis, ...]
+ imgs = np.ascontiguousarray(imgs, dtype=np.float32)
+
+ lb_affs = np.ascontiguousarray(lb_affs, dtype=np.float32)
+
+ bounds = bounds[np.newaxis, ...]
+ bounds = np.ascontiguousarray(bounds, dtype=np.float32)
+
+ bmask = bmask[np.newaxis, ...]
+ bmask = np.ascontiguousarray(bmask, dtype=np.float32)
+
+ return imgs, lb_affs, bounds, bmask
+
+ def perturbations_init(self):
+ self.per_intensity = Intensity()
+ self.per_gaussnoise = GaussNoise(min_std=self.min_noise_std, max_std=self.max_noise_std, norm_mode='trunc')
+ self.per_gaussblur = GaussBlur(min_kernel=self.min_kernel_size, max_kernel=self.max_kernel_size,
+ min_sigma=self.min_sigma, max_sigma=self.max_sigma)
+ self.per_cutout = Cutout(model_type=self.model_type)
+ self.per_misalign = Elastic(control_point_spacing=[4, 40, 40], jitter_sigma=[0, 0, 0], prob_slip=0.2,
+ prob_shift=0.2, max_misalign=17, padding=20)
+ self.per_elastic = Elastic(control_point_spacing=[4, 40, 40], jitter_sigma=[0, 2, 2], padding=20)
+
+ def apply_perturbations(self, data, mask):
+ if random.random() < 0.25:
+ data = self.per_intensity(data)
+ if random.random() < 0.25:
+ data = self.per_gaussnoise(data)
+ if random.random() < 0.25:
+ data = self.per_gaussblur(data)
+ if random.random() < 0.25:
+ data = self.per_cutout(data)
+ if random.random() < 0.25:
+ data, mask = self.per_elastic(data, mask)
+
+ return data, mask
+
+ def __len__(self):
+ return int(sys.maxsize)
+
+
+class Provider(object):
+ def __init__(self, stage, cfg):
+ self.stage = stage
+ if self.stage == 'train':
+ self.data = Train(cfg)
+ self.batch_size = cfg.TRAIN.batch_size
+ self.num_workers = cfg.TRAIN.num_workers
+ elif self.stage == 'valid':
+ pass
+ else:
+ raise AttributeError('Stage must be train/valid')
+ self.is_cuda = cfg.TRAIN.if_cuda
+ self.data_iter = None
+ self.iteration = 0
+ self.epoch = 1
+
+ def __len__(self):
+ return self.data.num_per_epoch
+
+ def build(self):
+ if self.stage == 'train':
+ self.data_iter = iter(
+ DataLoader(dataset=self.data, batch_size=self.batch_size, num_workers=self.num_workers,
+ shuffle=False, drop_last=False, pin_memory=True))
+ else:
+ self.data_iter = iter(DataLoader(dataset=self.data, batch_size=1, num_workers=0,
+ shuffle=False, drop_last=False, pin_memory=True))
+
+ def next(self):
+ if self.data_iter is None:
+ self.build()
+ try:
+ batch = next(self.data_iter)
+ self.iteration += 1
+ if self.is_cuda:
+ batch[0] = batch[0].cuda()
+ batch[1] = batch[1].cuda()
+ batch[2] = batch[2].cuda()
+ batch[3] = batch[3].cuda()
+ return batch
+ except StopIteration:
+ self.epoch += 1
+ self.build()
+ self.iteration += 1
+ batch = next(self.data_iter)
+ if self.is_cuda:
+ batch[0] = batch[0].cuda()
+ batch[1] = batch[1].cuda()
+ batch[2] = batch[2].cuda()
+ batch[3] = batch[3].cuda()
+ return batch
+
+
+def FDA_source_to_target_np(src_img, trg_img, L=0.1):
+ # exchange magnitude
+ # input: src_img, trg_img
+
+ src_img_np = src_img # .cpu().numpy()
+ trg_img_np = trg_img # .cpu().numpy()
+
+ # get fft of both source and target
+ fft_src_np = np.fft.fft2(src_img_np, axes=(-2, -1))
+ fft_trg_np = np.fft.fft2(trg_img_np, axes=(-2, -1))
+
+ # extract amplitude and phase of both ffts
+ amp_src, pha_src = np.abs(fft_src_np), np.angle(fft_src_np)
+ amp_trg, pha_trg = np.abs(fft_trg_np), np.angle(fft_trg_np)
+
+ # mutate the amplitude part of source with target
+ amp_src_ = low_freq_mutate_np(amp_src, amp_trg, L=L)
+
+ # mutated fft of source
+ fft_src_ = amp_src_ * np.exp(1j * pha_src)
+
+ # get the mutated image
+ src_in_trg = np.fft.ifft2(fft_src_, axes=(-2, -1))
+ src_in_trg = np.real(src_in_trg)
+
+ return src_in_trg
+
+
+def low_freq_mutate_np(amp_src, amp_trg, L=0.1):
+ a_src = np.fft.fftshift(amp_src, axes=(-2, -1))
+ a_trg = np.fft.fftshift(amp_trg, axes=(-2, -1))
+
+ h, w = a_src.shape
+ b = (np.floor(np.amin((h, w)) * L)).astype(int)
+ c_h = np.floor(h / 2.0).astype(int)
+ c_w = np.floor(w / 2.0).astype(int)
+
+ h1 = c_h - b
+ h2 = c_h + b + 1
+ w1 = c_w - b
+ w2 = c_w + b + 1
+
+ a_src[h1:h2, w1:w2] = a_trg[h1:h2, w1:w2]
+ a_src = np.fft.ifftshift(a_src, axes=(-2, -1))
+ return a_src
+
+
+def FDA_source_to_target_np_3D(src_img, trg_img, L=0.01):
+ scr_out = []
+ for i in range(src_img.shape[0]):
+ scr_out.append(FDA_source_to_target_np(src_img[i], trg_img[i], L))
+ return np.array(scr_out)
diff --git a/legacy/Train_and_Inference/supervised_train.py b/legacy/Train_and_Inference/supervised_train.py
new file mode 100644
index 0000000..a486917
--- /dev/null
+++ b/legacy/Train_and_Inference/supervised_train.py
@@ -0,0 +1,266 @@
+from __future__ import absolute_import
+from __future__ import print_function
+from __future__ import division
+
+import os
+
+os.environ['CUDA_VISIBLE_DEVICES'] = "0, 1"
+
+import sys
+import yaml
+import time
+import logging
+import argparse
+import numpy as np
+from attrdict import AttrDict
+from tensorboardX import SummaryWriter
+from collections import OrderedDict
+import torch
+import torch.nn as nn
+from supervised_provider import Provider
+from utils.show import show_affs
+from model.Mnet import MNet
+from utils.utils import setup_seed
+from loss.loss import WeightedBCE, WeightedMSE
+
+
+def init_project(cfg):
+ def init_logging(path):
+ logging.basicConfig(
+ level=logging.INFO,
+ format='%(message)s',
+ datefmt='%m-%d %H:%M',
+ filename=path,
+ filemode='w')
+ console = logging.StreamHandler()
+ console.setLevel(logging.INFO)
+ formatter = logging.Formatter('%(message)s')
+ console.setFormatter(formatter)
+ logging.getLogger('').addHandler(console)
+
+ setup_seed(cfg.TRAIN.random_seed)
+ if cfg.TRAIN.if_cuda:
+ if torch.cuda.is_available() is False:
+ raise AttributeError('No GPU available')
+
+ prefix = cfg.time
+ if cfg.TRAIN.resume:
+ model_name = cfg.TRAIN.model_name
+ else:
+ model_name = prefix + '_' + cfg.NAME
+ cfg.cache_path = os.path.join(cfg.TRAIN.cache_path, model_name)
+ cfg.save_path = os.path.join(cfg.TRAIN.save_path, model_name)
+ cfg.record_path = os.path.join(cfg.save_path, model_name)
+ cfg.valid_path = os.path.join(cfg.save_path, 'valid')
+ if cfg.TRAIN.resume is False:
+ if not os.path.exists(cfg.cache_path):
+ os.makedirs(cfg.cache_path)
+ if not os.path.exists(cfg.save_path):
+ os.makedirs(cfg.save_path)
+ if not os.path.exists(cfg.record_path):
+ os.makedirs(cfg.record_path)
+ if not os.path.exists(cfg.valid_path):
+ os.makedirs(cfg.valid_path)
+ init_logging(os.path.join(cfg.record_path, prefix + '.log'))
+ logging.info(cfg)
+ writer = SummaryWriter(cfg.record_path)
+ writer.add_text('cfg', str(cfg))
+ return writer
+
+
+def load_dataset(cfg):
+ print('Caching datasets ... ', flush=True)
+ t1 = time.time()
+ train_provider = Provider('train', cfg)
+ print('Done (time: %.2fs)' % (time.time() - t1))
+ return train_provider, valid_provider
+
+
+def build_model(cfg, writer):
+ print('Building model on ', end='', flush=True)
+ t1 = time.time()
+ device = torch.device('cuda:0')
+
+ print('load mnet!')
+
+ model = MNet(1, kn=(32, 64, 96, 128, 256), FMU='sub').cuda()
+
+ if cfg.MODEL.pre_train:
+ ckpt_path = cfg.MODEL.pretrain_path
+ print('Load pre-trained model from' + ckpt_path)
+ checkpoint = torch.load(ckpt_path)
+ pretrained_dict = OrderedDict()
+ state_dict = checkpoint['model_weights']
+ for k, v in state_dict.items():
+ name = k.replace('module.', '') if 'module' in k else k
+ pretrained_dict[name] = v
+
+ model_dict = model.state_dict()
+
+ pretrained_dict = {k: v for k, v in pretrained_dict.items() if
+ (k in model_dict) and ('up' not in k) and (
+ 'outputs' not in k)} # 1. filter out unnecessary keys
+ print(pretrained_dict)
+ model_dict.update(pretrained_dict) # 2. overwrite entries in the existing state dict
+ model.load_state_dict(model_dict)
+
+ if cfg.MODEL.continue_train:
+ ckpt_path = cfg.MODEL.continue_path
+ print('Load pre-trained model from' + ckpt_path)
+ checkpoint = torch.load(ckpt_path)
+ new_state_dict = OrderedDict()
+ state_dict = checkpoint['model_weights']
+ for k, v in state_dict.items():
+ name = k.replace('module.', '') if 'module' in k else k
+ new_state_dict[name] = v
+
+ model.load_state_dict(new_state_dict)
+ model = model.to(device)
+
+ cuda_count = torch.cuda.device_count()
+ if cuda_count > 1:
+ if cfg.TRAIN.batch_size % cuda_count == 0:
+ print('%d GPUs ... ' % cuda_count, end='', flush=True)
+ model = nn.DataParallel(model)
+ else:
+ raise AttributeError(
+ 'Batch size (%d) cannot be equally divided by GPU number (%d)' % (cfg.TRAIN.batch_size, cuda_count))
+ else:
+ print('a single GPU ... ', end='', flush=True)
+ print('Done (time: %.2fs)' % (time.time() - t1))
+ return model
+
+
+def calculate_lr(iters):
+ if iters < cfg.TRAIN.warmup_iters:
+ current_lr = (cfg.TRAIN.base_lr - cfg.TRAIN.end_lr) * pow(float(iters) / cfg.TRAIN.warmup_iters,
+ cfg.TRAIN.power) + cfg.TRAIN.end_lr
+ else:
+ if iters < cfg.TRAIN.decay_iters:
+ current_lr = (cfg.TRAIN.base_lr - cfg.TRAIN.end_lr) * pow(
+ 1 - float(iters - cfg.TRAIN.warmup_iters) / cfg.TRAIN.decay_iters, cfg.TRAIN.power) + cfg.TRAIN.end_lr
+ else:
+ current_lr = cfg.TRAIN.end_lr
+ return current_lr
+
+
+def loop(cfg, train_provider, model, optimizer, iters, writer):
+ f_loss_txt = open(os.path.join(cfg.record_path, 'loss.txt'), 'a')
+ rcd_time = []
+ sum_time = 0
+ sum_loss = 0
+ sum_labeled_loss = 0
+
+ if cfg.TRAIN.loss_func == 'MSELoss':
+ criterion = WeightedMSE()
+ elif cfg.TRAIN.loss_func == 'BCELoss':
+ criterion = WeightedBCE()
+ else:
+ raise AttributeError("NO this criterion")
+
+ while iters <= cfg.TRAIN.total_iters:
+ # train
+ model.train()
+ iters += 1
+ t1 = time.time()
+ inputs, target, bound, bmask = train_provider.next()
+
+ # decay learning rate
+ if cfg.TRAIN.end_lr == cfg.TRAIN.base_lr:
+ current_lr = cfg.TRAIN.base_lr
+ else:
+ current_lr = calculate_lr(iters)
+ for param_group in optimizer.param_groups:
+ param_group['lr'] = current_lr
+
+ optimizer.zero_grad()
+ pred, pred_b = model(inputs)
+
+ ##############################
+ loss1 = criterion(pred, target, weight=bmask)
+ loss2 = criterion(pred_b, bound, weight=bmask)
+ loss = loss1 + loss2
+ loss.backward()
+ ##############################
+
+ if cfg.TRAIN.weight_decay is not None:
+ for group in optimizer.param_groups:
+ for param in group['params']:
+ param.data = param.data.add(-cfg.TRAIN.weight_decay * group['lr'], param.data)
+ optimizer.step()
+
+ sum_loss += loss.item()
+ sum_time += time.time() - t1
+
+ # log train
+ if iters % cfg.TRAIN.display_freq == 0 or iters == 1:
+ rcd_time.append(sum_time)
+ if iters == 1:
+ logging.info(
+ 'step %d, loss = %.6f, labeled_loss=%.6f (wt: *1, lr: %.8f, et: %.2f sec, rd: %.2f min)'
+ % (iters, sum_loss, sum_labeled_loss, current_lr, sum_time,
+ (cfg.TRAIN.total_iters - iters) / cfg.TRAIN.display_freq * np.mean(np.asarray(rcd_time)) / 60))
+ writer.add_scalar('loss', sum_loss * 1, iters)
+ else:
+ logging.info(
+ 'step %d, loss = %.6f, labeled_loss=%.6f (wt: *1, lr: %.8f, et: %.2f sec, rd: %.2f min)' \
+ % (iters, sum_loss / cfg.TRAIN.display_freq * 1, \
+ sum_labeled_loss / cfg.TRAIN.display_freq * 1, \
+ current_lr, sum_time, \
+ (cfg.TRAIN.total_iters - iters) / cfg.TRAIN.display_freq * np.mean(np.asarray(rcd_time)) / 60))
+ writer.add_scalar('loss', sum_loss / cfg.TRAIN.display_freq * 1, iters)
+ f_loss_txt.write('step = %d, loss = %.6f, labeled_loss=%.6f' \
+ % (iters, sum_loss / cfg.TRAIN.display_freq * 1, \
+ sum_labeled_loss / cfg.TRAIN.display_freq * 1))
+ f_loss_txt.write('\n')
+ f_loss_txt.flush()
+ sys.stdout.flush()
+ sum_time = 0
+ sum_loss = 0
+
+ # display
+ if iters % cfg.TRAIN.valid_freq == 0 or iters == 1:
+ show_affs(iters, inputs, pred, target, cfg.cache_path, model_type=cfg.MODEL.model_type)
+
+ # save
+ if iters % cfg.TRAIN.save_freq == 0 and iters >= 0:
+ states = {'current_iter': iters, 'valid_result': None,
+ 'model_weights': model.state_dict()}
+ torch.save(states, os.path.join(cfg.save_path, 'model-%06d.ckpt' % iters))
+ print('***************save modol, iters = %d.***************' % (iters), flush=True)
+
+ f_loss_txt.close()
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument('-c', '--cfg', type=str, default='SegNeuron', help='path to config file')
+ parser.add_argument('-m', '--mode', type=str, default='train', help='path to config file')
+ args = parser.parse_args()
+
+ cfg_file = args.cfg + '.yaml'
+ print('cfg_file: ' + cfg_file)
+ print('mode: ' + args.mode)
+
+ with open('./config/' + cfg_file, 'r') as f:
+ cfg = AttrDict(yaml.safe_load(f))
+
+ timeArray = time.localtime()
+ time_stamp = time.strftime('%Y-%m-%d--%H-%M-%S', timeArray)
+ print('time stamp:', time_stamp)
+
+ cfg.path = cfg_file
+ cfg.time = time_stamp
+
+ if args.mode == 'train':
+ writer = init_project(cfg)
+ train_provider, valid_provider = load_dataset(cfg)
+ model = build_model(cfg, writer)
+ optimizer = torch.optim.Adam(model.parameters(), lr=cfg.TRAIN.base_lr, betas=(0.9, 0.999),
+ eps=0.01, weight_decay=1e-6, amsgrad=True)
+ init_iters = 0
+ loop(cfg, train_provider, valid_provider, model, optimizer, init_iters, writer)
+ writer.close()
+ else:
+ pass
+ print('***Done***')
diff --git a/legacy/Train_and_Inference/utils/aff_util.py b/legacy/Train_and_Inference/utils/aff_util.py
new file mode 100644
index 0000000..322251a
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/aff_util.py
@@ -0,0 +1,139 @@
+import numpy as np
+#from em_segLib.seg_util import check_volume
+# from scipy.misc import comb
+from scipy.special import comb
+import scipy.sparse
+
+
+def affinitize(img, ret=None, dst=(1,1,1), dtype='float32'):
+ # PNI code
+ """
+ Transform segmentation to an affinity map.
+ Args:
+ img: 3D indexed image, with each index corresponding to each segment.
+ Returns:
+ ret: an affinity map (4D tensor).
+ """
+ img = check_volume(img)
+ if ret is None:
+ ret = np.zeros(img.shape, dtype=dtype)
+
+ # Sanity check.
+ (dz,dy,dx) = dst
+ assert abs(dx) < img.shape[-1]
+ assert abs(dy) < img.shape[-2]
+ assert abs(dz) < img.shape[-3]
+
+ # Slices.
+ s0 = list()
+ s1 = list()
+ s2 = list()
+ for i in range(3):
+ if dst[i] == 0:
+ s0.append(slice(None))
+ s1.append(slice(None))
+ s2.append(slice(None))
+ elif dst[i] > 0:
+ s0.append(slice(dst[i], None))
+ s1.append(slice(dst[i], None))
+ s2.append(slice(None, -dst[i]))
+ else:
+ s0.append(slice(None, dst[i]))
+ s1.append(slice(-dst[i], None))
+ s2.append(slice(None, dst[i]))
+
+ ret[s0] = (img[s1]==img[s2]) & (img[s1]>0)
+ return ret[np.newaxis,...]
+
+def bmap_to_affgraph(bmap,nhood,return_min_idx=False):
+ # constructs an affinity graph from a boundary map
+ # assume affinity graph is represented as:
+ # shape = (e, z, y, x)
+ # nhood.shape = (edges, 3)
+ shape = bmap.shape
+ nEdge = nhood.shape[0]
+ aff = np.zeros((nEdge,)+shape,dtype=np.int32)
+ minidx = np.zeros((nEdge,)+shape,dtype=np.int32)
+
+ for e in range(nEdge):
+ aff[e, \
+ max(0,-nhood[e,0]):min(shape[0],shape[0]-nhood[e,0]), \
+ max(0,-nhood[e,1]):min(shape[1],shape[1]-nhood[e,1]), \
+ max(0,-nhood[e,2]):min(shape[2],shape[2]-nhood[e,2])] = np.minimum( \
+ bmap[max(0,-nhood[e,0]):min(shape[0],shape[0]-nhood[e,0]), \
+ max(0,-nhood[e,1]):min(shape[1],shape[1]-nhood[e,1]), \
+ max(0,-nhood[e,2]):min(shape[2],shape[2]-nhood[e,2])], \
+ bmap[max(0,nhood[e,0]):min(shape[0],shape[0]+nhood[e,0]), \
+ max(0,nhood[e,1]):min(shape[1],shape[1]+nhood[e,1]), \
+ max(0,nhood[e,2]):min(shape[2],shape[2]+nhood[e,2])] )
+ minidx[e, \
+ max(0,-nhood[e,0]):min(shape[0],shape[0]-nhood[e,0]), \
+ max(0,-nhood[e,1]):min(shape[1],shape[1]-nhood[e,1]), \
+ max(0,-nhood[e,2]):min(shape[2],shape[2]-nhood[e,2])] = \
+ bmap[max(0,-nhood[e,0]):min(shape[0],shape[0]-nhood[e,0]), \
+ max(0,-nhood[e,1]):min(shape[1],shape[1]-nhood[e,1]), \
+ max(0,-nhood[e,2]):min(shape[2],shape[2]-nhood[e,2])] > \
+ bmap[max(0,nhood[e,0]):min(shape[0],shape[0]+nhood[e,0]), \
+ max(0,nhood[e,1]):min(shape[1],shape[1]+nhood[e,1]), \
+ max(0,nhood[e,2]):min(shape[2],shape[2]+nhood[e,2])]
+ return aff
+
+def seg_to_affgraph(seg,nhood,pad=''):
+ # constructs an affinity graph from a segmentation
+ # assume affinity graph is represented as:
+ # shape = (e, z, y, x)
+ # nhood.shape = (edges, 3)
+ shape = seg.shape
+ nEdge = nhood.shape[0]
+ aff = np.zeros((nEdge,)+shape,dtype=np.int32)
+
+ for e in range(nEdge):
+ aff[e, \
+ max(0,-nhood[e,0]):min(shape[0],shape[0]-nhood[e,0]), \
+ max(0,-nhood[e,1]):min(shape[1],shape[1]-nhood[e,1]), \
+ max(0,-nhood[e,2]):min(shape[2],shape[2]-nhood[e,2])] = \
+ (seg[max(0,-nhood[e,0]):min(shape[0],shape[0]-nhood[e,0]), \
+ max(0,-nhood[e,1]):min(shape[1],shape[1]-nhood[e,1]), \
+ max(0,-nhood[e,2]):min(shape[2],shape[2]-nhood[e,2])] == \
+ seg[max(0,nhood[e,0]):min(shape[0],shape[0]+nhood[e,0]), \
+ max(0,nhood[e,1]):min(shape[1],shape[1]+nhood[e,1]), \
+ max(0,nhood[e,2]):min(shape[2],shape[2]+nhood[e,2])] ) \
+ * ( seg[max(0,-nhood[e,0]):min(shape[0],shape[0]-nhood[e,0]), \
+ max(0,-nhood[e,1]):min(shape[1],shape[1]-nhood[e,1]), \
+ max(0,-nhood[e,2]):min(shape[2],shape[2]-nhood[e,2])] > 0 ) \
+ * ( seg[max(0,nhood[e,0]):min(shape[0],shape[0]+nhood[e,0]), \
+ max(0,nhood[e,1]):min(shape[1],shape[1]+nhood[e,1]), \
+ max(0,nhood[e,2]):min(shape[2],shape[2]+nhood[e,2])] > 0 )
+ if nEdge==3 and pad == 'replicate': # pad the boundary affinity
+ aff[0,0] = (seg[0]>0).astype(aff.dtype)
+ aff[1,:,0] = (seg[:,0]>0).astype(aff.dtype)
+ aff[2,:,:,0] = (seg[:,:,0]>0).astype(aff.dtype)
+
+ return aff
+
+def affgraph_to_edgelist(aff,nhood):
+ node1,node2 = nodelist_like(aff.shape[1:],nhood)
+ return (node1.ravel(),node2.ravel(),aff.ravel())
+
+def nodelist_like(shape,nhood):
+ # constructs the node lists corresponding to the edge list representation of an affinity graph
+ # assume node shape is represented as:
+ # shape = (z, y, x)
+ # nhood.shape = (edges, 3)
+ nEdge = nhood.shape[0]
+ nodes = np.arange(np.prod(shape),dtype=np.uint64).reshape(shape)
+ node1 = np.tile(nodes,(nEdge,1,1,1))
+ node2 = np.full(node1.shape,-1,dtype=np.uint64)
+
+ for e in range(nEdge):
+ node2[e, \
+ max(0,-nhood[e,0]):min(shape[0],shape[0]-nhood[e,0]), \
+ max(0,-nhood[e,1]):min(shape[1],shape[1]-nhood[e,1]), \
+ max(0,-nhood[e,2]):min(shape[2],shape[2]-nhood[e,2])] = \
+ nodes[max(0,nhood[e,0]):min(shape[0],shape[0]+nhood[e,0]), \
+ max(0,nhood[e,1]):min(shape[1],shape[1]+nhood[e,1]), \
+ max(0,nhood[e,2]):min(shape[2],shape[2]+nhood[e,2])]
+
+ return (node1, node2)
+
+
diff --git a/legacy/Train_and_Inference/utils/affine.py b/legacy/Train_and_Inference/utils/affine.py
new file mode 100644
index 0000000..4e63b8c
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/affine.py
@@ -0,0 +1,288 @@
+import numpy as np
+
+def identity_xf(N):
+ """
+ Construct N identity 2x3 transformation matrices
+ :return: array of shape (N, 2, 3)
+ """
+ # Taken from https://github.com/Britefury/pytorch-mask-rcnn/blob/refactor/maskrcnn/utils/affine_transforms.py
+ xf = np.zeros((N, 2, 3), dtype=np.float32)
+ xf[:, 0, 0] = xf[:, 1, 1] = 1.0
+ return xf
+
+
+def inv_nx2x2(X):
+ """
+ Invert the N 2x2 transformation matrices stored in X; a (N,2,2) array
+ :param X: transformation matrices to invert, (N,2,2) array
+ :return: inverse of X
+ """
+ # Taken from https://github.com/Britefury/pytorch-mask-rcnn/blob/refactor/maskrcnn/utils/affine_transforms.py
+ rdet = 1.0 / (X[:, 0, 0] * X[:, 1, 1] - X[:, 1, 0] * X[:, 0, 1])
+ y = np.zeros_like(X)
+ y[:, 0, 0] = X[:, 1, 1] * rdet
+ y[:, 1, 1] = X[:, 0, 0] * rdet
+ y[:, 0, 1] = -X[:, 0, 1] * rdet
+ y[:, 1, 0] = -X[:, 1, 0] * rdet
+ return y
+
+def inv_nx2x3(m):
+ """
+ Invert the N 2x3 transformation matrices stored in X; a (N,2,3) array
+ :param X: transformation matrices to invert, (N,2,3) array
+ :return: inverse of X
+ """
+ # Taken from https://github.com/Britefury/pytorch-mask-rcnn/blob/refactor/maskrcnn/utils/affine_transforms.py
+ m2 = m[:, :, :2]
+ mx = m[:, :, 2:3]
+ m2inv = inv_nx2x2(m2)
+ mxinv = np.matmul(m2inv, -mx)
+ return np.append(m2inv, mxinv, axis=2)
+
+def cat_nx2x3_2(a, b):
+ """
+ Multiply the N 2x3 transformations stored in `a` with those in `b`
+ :param a: transformation matrices, (N,2,3) array
+ :param b: transformation matrices, (N,2,3) array
+ :return: `a . b`
+ """
+ # Taken from https://github.com/Britefury/pytorch-mask-rcnn/blob/refactor/maskrcnn/utils/affine_transforms.py
+ a2 = a[:, :, :2]
+ b2 = b[:, :, :2]
+
+ ax = a[:, :, 2:3]
+ bx = b[:, :, 2:3]
+
+ ab2 = np.matmul(a2, b2)
+ abx = ax + np.matmul(a2, bx)
+ return np.append(ab2, abx, axis=2)
+
+def cat_nx2x3(*x):
+ """
+ Multiply the N 2x3 transformations stored in the arrays in `x`
+ :param x: transformation matrices, tuple of (N,2,3) arrays
+ :return: `x[0] . x[1] . ... . x[N-1]`
+ """
+ # Taken from https://github.com/Britefury/pytorch-mask-rcnn/blob/refactor/maskrcnn/utils/affine_transforms.py
+ y = x[0]
+ for i in range(1, len(x)):
+ y = cat_nx2x3_2(y, x[i])
+ return y
+
+def translation_matrices(xlats_xy):
+ """
+ Generate translation matrices
+ :param xlats_xy: translations as an (N, 2) array (x,y)
+ :return: translations matrices, (N,2,3) array
+ """
+ # Taken from https://github.com/Britefury/pytorch-mask-rcnn/blob/refactor/maskrcnn/utils/affine_transforms.py
+ N = len(xlats_xy)
+ xf = np.zeros((N, 2, 3), dtype=np.float32)
+ xf[:, 0, 0] = xf[:, 1, 1] = 1.0
+ xf[:, :, 2] = xlats_xy
+ return xf
+
+def scale_matrices(scale_xy):
+ """
+ Generate translation matrices
+ :param scale_xy: scale factors as an (N, 2) array (x,y)
+ :return: translations matrices, (N,2,3) array
+ """
+ # Taken from https://github.com/Britefury/pytorch-mask-rcnn/blob/refactor/maskrcnn/utils/affine_transforms.py
+ N = len(scale_xy)
+ xf = np.zeros((N, 2, 3), dtype=np.float32)
+ xf[:, 0, 0] = scale_xy[:, 0]
+ xf[:, 1, 1] = scale_xy[:, 1]
+ return xf
+
+def rotation_matrices(thetas):
+ """
+ Generate rotation matrices
+
+ Counter-clockwise, +y points downwards
+
+ Where s = sin(theta) and c = cos(theta)
+
+ M = [[ c s 0 ]
+ [ -s c 0 ]]
+
+ :param thetas: rotation angles in radians as a (N,) array
+ :return: rotation matrices, (N,2,3) array
+ """
+ # Taken from https://github.com/Britefury/pytorch-mask-rcnn/blob/refactor/maskrcnn/utils/affine_transforms.py
+ N = len(thetas)
+ c = np.cos(thetas)
+ s = np.sin(thetas)
+ rot_xf = np.zeros((N, 2, 3), dtype=np.float32)
+ rot_xf[:, 0, 0] = rot_xf[:, 1, 1] = c
+ rot_xf[:, 1, 0] = -s
+ rot_xf[:, 0, 1] = s
+ return rot_xf
+
+def flip_xyd_matrices(flip_flags_xyd, image_size):
+ """
+ Generate flip matrices in OpenCV compatible form. Each sample has three flags: `x`, `y` and `d`:
+ `x == True` -> flip horizontally
+ `y == True` -> flip vertically
+ `d == True` -> flip diagonal or swap X and Y axes
+
+ :param flip_flags_xyd: per sample flip flags as a (N,[x, y, d]) array
+ :param image_size: image size as a `(H, w)` tuple
+ :return: flip matrices, (N,2,3) array
+ """
+ if flip_flags_xyd.ndim != 2:
+ raise ValueError('flip_flags_xyd should have 2 dimensions, not {}'.format(flip_flags_xyd.ndim))
+ if flip_flags_xyd.shape[1] != 3:
+ raise ValueError('flip_flags_xyd.shape[1] should be 3 dimensions, not {}'.format(flip_flags_xyd.shape[1]))
+
+ # False -> 1, True -> -1
+ flip_scale_xy = flip_flags_xyd[:, :2] * -2 + 1
+ # Negative scale factors need to be combined with a translation whose value is (image_size - 1)
+ # Mask the translation with the flip flags to only apply it where flipping is done
+ flip_xlat_xy = flip_flags_xyd[:, :2] * (np.array(image_size[::-1]).astype(float) - 1)
+
+ hv_flip_xf = identity_xf(len(flip_flags_xyd))
+
+ # Diagonal flip: swap X and Y axes
+ diag = flip_flags_xyd[:, 2]
+ hv_flip_xf[diag] = hv_flip_xf[diag][:, ::-1, :]
+
+ return cat_nx2x3(
+ hv_flip_xf,
+ translation_matrices(flip_xlat_xy),
+ scale_matrices(flip_scale_xy),
+ )
+
+
+
+def centre_xf(xf, size):
+ """
+ Centre the transformations in `xf` around (0,0), where the current centre is assumed to be at the
+ centre of an image of shape `size`
+ :param xf: transformation matrices, (N,2,3) array
+ :param size: image size
+ :return: centred transformation matrices, (N,2,3) array
+ """
+ # Taken from https://github.com/Britefury/pytorch-mask-rcnn/blob/refactor/maskrcnn/utils/affine_transforms.py
+ height, width = size
+
+ # centre_to_zero moves the centre of the image to (0,0)
+ centre_to_zero = np.zeros((1, 2, 3), dtype=np.float32)
+ centre_to_zero[0, 0, 0] = centre_to_zero[0, 1, 1] = 1.0
+ centre_to_zero[0, 0, 2] = -float(width) * 0.5
+ centre_to_zero[0, 1, 2] = -float(height) * 0.5
+
+ # centre_to_zero then xf
+ xf_centred = cat_nx2x3(xf, centre_to_zero)
+
+ # move (0,0) back to the centre
+ xf_centred[:, 0, 2] += float(width) * 0.5
+ xf_centred[:, 1, 2] += float(height) * 0.5
+
+ return xf_centred
+
+
+def cv_to_torch(mtx, dst_size, src_size=None):
+ """
+ Convert transformations matrices that can be used with `cv2.warpAffine` to work with PyTorch
+ grid sampling.
+
+ NOTE: `align_corners=True` should be passed to `F.affine_Grid` and `F.grid_sample` to
+ correctly match OpenCV transformations.
+
+ `cv2.warpAffine` expects a matrix that transforms an image in pixel co-ordinates.
+ PyTorch `F.affine_grid` and `F.grid_sample` maps pixel locations to a [-1, 1] grid
+ and transforms these sample locations, prior to sampling the image.
+
+ :param mtx: OpenCV transformation matrices as a (N,2,3) array
+ :param dst_size: the size of the output image as a `(height, width)` tuple
+ :param src_size: the size of the input image as a `(height, width)` tuple, or None to use `dst_size`
+ :return: PyTorch transformation matrices as a (N,2,3) array
+ """
+ # Taken from https://github.com/Britefury/pytorch-mask-rcnn/blob/refactor/maskrcnn/utils/affine_transforms.py
+ dst_scale_x = float(dst_size[1] - 1) / 2.0
+ dst_scale_y = float(dst_size[0] - 1) / 2.0
+
+ if src_size is not None:
+ src_scale_x = float(src_size[1] - 1) / 2.0
+ src_scale_y = float(src_size[0] - 1) / 2.0
+ else:
+ src_scale_x = dst_scale_x
+ src_scale_y = dst_scale_y
+
+ N = len(mtx)
+
+ # OpenCV transforms the image, whereas the PyTorch transforms the points at which the
+ # image is samples. We account for this by inverting the transformation matrices
+ mtx = inv_nx2x3(mtx)
+
+ torch_cv = identity_xf(N)
+ torch_cv[:, 0, 0] = dst_scale_x
+ torch_cv[:, 1, 1] = dst_scale_y
+ torch_cv[:, 0, 2] = dst_scale_x
+ torch_cv[:, 1, 2] = dst_scale_y
+
+ cv_torch = identity_xf(N)
+ cv_torch[:, 0, 0] = 1.0 / src_scale_x
+ cv_torch[:, 1, 1] = 1.0 / src_scale_y
+ cv_torch[:, 0, 2] = -1.0
+ cv_torch[:, 1, 2] = -1.0
+
+ # Transform torch co-ordinates to OpenCV, apply the transformation then OpenCV co-ordinates back to torch
+ return cat_nx2x3(cv_torch, mtx, torch_cv)
+
+
+def pil_to_torch(mtx, dst_size, src_size=None, align_corners=True):
+ """
+ Convert affine transformations matrices that can be used with Pillow `Image.transform` to work with PyTorch
+ grid sampling.
+
+ `Image.transform` expects a matrix that transforms an image in pixel co-ordinates, where pixel [0,0]
+ is centred at [0.5, 0.5].
+ PyTorch `F.affine_grid` and `F.grid_sample` maps pixel locations to a [-1, 1] grid
+ and transforms these sample locations, prior to sampling the image.
+
+ :param mtx: PIL transformation matrices as a (N,2,3) array
+ :param dst_size: the size of the output image as a `(height, width)` tuple
+ :param src_size: the size of the input image as a `(height, width)` tuple, or None to use `dst_size`
+ :param align_corners: if you want to use `align_corners=False` for PyTorch `F.affine_grid` and `F.grid_sample`,
+ pass `align_corners=False` here
+ :return: PyTorch transformation matrices as a (N,2,3) array
+ """
+ # Taken from https://github.com/Britefury/pytorch-mask-rcnn/blob/refactor/maskrcnn/utils/affine_transforms.py
+ if align_corners:
+ dst_size = (dst_size[0] - 1, dst_size[1] - 1)
+ dst_scale_x = float(dst_size[1]) / 2.0
+ dst_scale_y = float(dst_size[0]) / 2.0
+
+ if src_size is not None:
+ if align_corners:
+ src_size = (src_size[0] - 1, src_size[1] - 1)
+ src_scale_x = float(src_size[1]) / 2.0
+ src_scale_y = float(src_size[0]) / 2.0
+ else:
+ src_scale_x = dst_scale_x
+ src_scale_y = dst_scale_y
+
+ N = len(mtx)
+
+ torch_cv = identity_xf(N)
+ torch_cv[:, 0, 0] = dst_scale_x
+ torch_cv[:, 1, 1] = dst_scale_y
+ torch_cv[:, 0, 2] = dst_scale_x
+ torch_cv[:, 1, 2] = dst_scale_y
+ if align_corners:
+ torch_cv[:, 0, 2] += 0.5
+ torch_cv[:, 1, 2] += 0.5
+
+ cv_torch = identity_xf(N)
+ cv_torch[:, 0, 0] = 1.0 / src_scale_x
+ cv_torch[:, 1, 1] = 1.0 / src_scale_y
+ cv_torch[:, 0, 2] = -1.0
+ cv_torch[:, 1, 2] = -1.0
+ if align_corners:
+ cv_torch[:, 0, 2] += -0.5 / src_scale_x
+ cv_torch[:, 1, 2] += -0.5 / src_scale_y
+
+ # Transform torch co-ordinates to OpenCV, apply the transformation then OpenCV co-ordinates back to torch
+ return cat_nx2x3(cv_torch, mtx, torch_cv)
diff --git a/legacy/Train_and_Inference/utils/augmentation.py b/legacy/Train_and_Inference/utils/augmentation.py
new file mode 100644
index 0000000..e0dcd95
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/augmentation.py
@@ -0,0 +1,737 @@
+## FUNCTIONS: data augmentation used for "superhuman" network
+## Noted that it is only applied to 3D datasets
+## Written by Wei Huang
+## 2020/10/12
+## reference: https://github.com/donglaiw/EM-network/blob/master/em_net/data/augmentation.py
+
+import cv2
+import time
+import math
+import random
+import torch
+import numpy as np
+from scipy.ndimage.interpolation import map_coordinates, zoom
+from scipy.ndimage.filters import gaussian_filter
+
+from utils.coordinate import Coordinate
+
+def produce_simple_aug(data, rule):
+ '''Routine data augmentation, including flipping in x-, y- and z-dimensions,
+ and transposing x- and y-dimensions, they have 2^4=16 combinations
+ Args:
+ data: numpy array, [Z, Y, X], ndim=3
+ rule: numpy array, list or tuple, but len(rule) = 4, such as rule=[1,1,0,0]
+ '''
+ assert data.ndim == 3 and len(rule) == 4
+ # z reflection.
+ if rule[0]:
+ data = data[::-1, :, :]
+ # x reflection.
+ if rule[1]:
+ data = data[:, :, ::-1]
+ # y reflection.
+ if rule[2]:
+ data = data[:, ::-1, :]
+ # Transpose in xy.
+ if rule[3]:
+ data = data.transpose(0, 2, 1)
+ return data
+
+########################################################################
+def create_identity_transformation(shape, subsample=1):
+ dims = len(shape)
+ subsample_shape = tuple(max(1, int(s/subsample)) for s in shape)
+ step_width = tuple(float(shape[d]-1)/(subsample_shape[d]-1)
+ if subsample_shape[d] > 1 else 1 for d in range(dims))
+
+ axis_ranges = (
+ np.arange(subsample_shape[d], dtype=np.float32)*step_width[d]
+ for d in range(dims)
+ )
+ return np.array(np.meshgrid(*axis_ranges, indexing='ij'), dtype=np.float32)
+
+
+def upscale_transformation(transformation,
+ output_shape,
+ interpolate_order=1):
+ input_shape = transformation.shape[1:]
+ dims = len(output_shape)
+ scale = tuple(float(s)/c for s, c in zip(output_shape, input_shape))
+
+ scaled = np.zeros((dims,)+output_shape, dtype=np.float32)
+ for d in range(dims):
+ zoom(transformation[d], zoom=scale,
+ output=scaled[d], order=interpolate_order)
+ return scaled
+
+
+def create_elastic_transformation(shape,
+ control_point_spacing=100,
+ jitter_sigma=10.0,
+ subsample=1):
+ dims = len(shape)
+ subsample_shape = tuple(max(1, int(s/subsample)) for s in shape)
+
+ try:
+ spacing = tuple((d for d in control_point_spacing))
+ except:
+ spacing = (control_point_spacing,)*dims
+ try:
+ sigmas = [s for s in jitter_sigma]
+ except:
+ sigmas = [jitter_sigma]*dims
+
+ control_points = tuple(
+ max(1, int(round(float(shape[d])/spacing[d])))
+ for d in range(len(shape))
+ )
+
+ # jitter control points
+ control_point_offsets = np.zeros(
+ (dims,) + control_points, dtype=np.float32)
+ for d in range(dims):
+ if sigmas[d] > 0:
+ control_point_offsets[d] = np.random.normal(
+ scale=sigmas[d], size=control_points)
+ transform = upscale_transformation(control_point_offsets, subsample_shape, interpolate_order=3)
+ return transform
+
+
+def rotate(point, angle):
+ res = np.array(point)
+ res[0] = math.sin(angle)*point[1] + math.cos(angle)*point[0]
+ res[1] = -math.sin(angle)*point[0] + math.cos(angle)*point[1]
+ return res
+
+
+def create_rotation_transformation(shape, angle, subsample=1):
+ dims = len(shape)
+ subsample_shape = tuple(max(1, int(s/subsample)) for s in shape)
+ control_points = (2,)*dims
+
+ # map control points to world coordinates
+ control_point_scaling_factor = tuple(float(s-1) for s in shape)
+
+ # rotate control points
+ center = np.array([0.5*(d-1) for d in shape])
+
+ control_point_offsets = np.zeros(
+ (dims,) + control_points, dtype=np.float32)
+ for control_point in np.ndindex(control_points):
+ point = np.array(control_point)*control_point_scaling_factor
+ center_offset = np.array(
+ [p-c for c, p in zip(center, point)], dtype=np.float32)
+ rotated_offset = np.array(center_offset)
+ rotated_offset[-2:] = rotate(center_offset[-2:], angle)
+ displacement = rotated_offset - center_offset
+ control_point_offsets[(slice(None),) + control_point] += displacement
+ return upscale_transformation(control_point_offsets, subsample_shape)
+
+
+def random_offset(max_misalign):
+ return Coordinate((0,) + tuple(max_misalign - random.randint(0, 2*int(max_misalign)) for d in range(2)))
+
+
+def misalign(transformation, prob_slip, prob_shift, max_misalign):
+ num_sections = transformation[0].shape[0]
+ shifts = [Coordinate((0, 0, 0))]*num_sections
+ # orginal
+ # for z in range(num_sections):
+ # r = random.random()
+ # if r <= prob_slip:
+ # shifts[z] = random_offset(max_misalign)
+ # elif r <= prob_slip + prob_shift:
+ # offset = random_offset(max_misalign)
+ # for zp in range(z, num_sections):
+ # shifts[zp] += offset
+
+ # written by Wei Huang
+ if random.random() > 0.5:
+ # slip type
+ for z in range(1, num_sections):
+ if random.random() <= prob_slip:
+ shifts[z] = random_offset(max_misalign)
+ else:
+ # translation type
+ for z in range(1, num_sections):
+ if random.random() <= prob_shift:
+ offset = random_offset(max_misalign)
+ for zp in range(z, num_sections):
+ shifts[zp] = offset
+ break
+
+ for z in range(num_sections):
+ transformation[1][z, :, :] += shifts[z][1]
+ transformation[2][z, :, :] += shifts[z][2]
+ return transformation
+
+def apply_transformation(image,
+ transformation,
+ interpolate=True,
+ outside_value=0,
+ output=None):
+ order = 1 if interpolate == True else 0
+ output = image.dtype if output is None else output
+ return map_coordinates(image,
+ transformation,
+ output=output,
+ order=order,
+ mode='constant',
+ cval=outside_value)
+
+def resize_3d(imgs, det_size, mode='linear'):
+ new_imgs = []
+ for k in range(imgs.shape[0]):
+ temp = imgs[k]
+ if mode == 'linear':
+ temp = cv2.resize(temp, (det_size, det_size), interpolation=cv2.INTER_LINEAR)
+ elif mode == 'nearest':
+ temp = cv2.resize(temp, (det_size, det_size), interpolation=cv2.INTER_NEAREST)
+ else:
+ raise AttributeError('No this interpolation mode!')
+ new_imgs.append(temp)
+ new_imgs = np.asarray(new_imgs)
+ return new_imgs
+
+########################################################################
+
+class Rescale(object):
+ def __init__(self, scale_factor=2, det_shape=[18, 160, 160]):
+ super(Rescale, self).__init__()
+ self.scale_factor = scale_factor
+ self.det_shape = det_shape
+
+ def __call__(self, data, mask):
+ src_shape = data.shape
+ assert src_shape[-1] >= self.det_shape[-1] * self.scale_factor, 'data shape must be 160*2'
+ min_size = self.det_shape[-1] // self.scale_factor
+ max_size = self.det_shape[-1] * self.scale_factor
+ scale_size = random.randint(min_size // 2, max_size // 2)
+ scale_size = scale_size * 2
+
+ if scale_size < src_shape[-1]:
+ shift = (src_shape[-1] - scale_size) // 2
+ data = data[:, shift:-shift, shift:-shift]
+ data = resize_3d(data, self.det_shape[-1], mode='linear')
+ mask = resize_3d(mask, self.det_shape[-1], mode='nearest')
+ return data, mask, scale_size
+
+class SimpleAugment(object):
+ def __init__(self, skip_ratio=0.5):
+ '''Routine data augmentation, including flipping in x-, y- and z-dimensions,
+ and transposing x- and y-dimensions, they have 2^4=16 combinations
+ Args:
+ skip_ratio: Probability of execution
+ '''
+ super(SimpleAugment, self).__init__()
+ self.ratio = skip_ratio
+
+ def __call__(self, inputs):
+ return self.forward(inputs)
+
+ def forward(self, inputs):
+ '''
+ Args:
+ inputs: list, such as [imgs, label, ...], imgs and label are numpy arrays with ndim=3
+ '''
+ skiprand = np.random.rand()
+ if skiprand < self.ratio:
+ rule = np.random.randint(2, size=4)
+ for idx in range(len(inputs)):
+ inputs[idx] = produce_simple_aug(inputs[idx], rule)
+ return inputs
+ else:
+ return inputs
+
+
+class RandomRotationAugment(object):
+ def __init__(self, skip_ratio=0.5):
+ '''Random rotation augmentation in x-y plane
+ Args:
+ skip_ratio: Probability of execution
+ '''
+ super(RandomRotationAugment, self).__init__()
+ self.ratio = skip_ratio
+
+ def __call__(self, inputs, mask=None):
+ return self.forward(inputs, mask)
+
+ def forward(self, inputs, mask=None):
+ '''
+ Args:
+ inputs: list, such as [imgs, label, ...], imgs and label are numpy arrays with ndim=3
+ '''
+ skiprand = np.random.rand()
+ if skiprand < self.ratio:
+ angle = random.randint(0, 360-1)
+ center = tuple(np.array(inputs.shape)[1:] // 2)
+ rot_mat = cv2.getRotationMatrix2D(center, angle, 1)
+ for k in range(inputs.shape[0]):
+ inputs[k] = cv2.warpAffine(inputs[k], rot_mat, inputs[k].shape, flags=cv2.INTER_LINEAR)
+ if mask is not None:
+ for k in range(mask.shape[0]):
+ mask[k] = cv2.warpAffine(mask[k], rot_mat, mask[k].shape, flags=cv2.INTER_NEAREST)
+ return inputs, mask
+ else:
+ return inputs
+ else:
+ if mask is not None:
+ return inputs, mask
+ else:
+ return inputs
+
+class IntensityAugment(object):
+ def __init__(self, mode='mix',
+ skip_ratio=0.5,
+ CONTRAST_FACTOR=0.1,
+ BRIGHTNESS_FACTOR=0.1):
+ '''Image intensity augmentation, including adjusting contrast and brightness
+ Args:
+ mode: '2D', '3D' or 'mix' (contains '2D' and '3D')
+ skip_ratio: Probability of execution
+ CONTRAST_FACTOR: Contrast factor
+ BRIGHTNESS_FACTOR : Brightness factor
+ '''
+ super(IntensityAugment, self).__init__()
+ assert mode == '3D' or mode == '2D' or mode == 'mix'
+ self.mode = mode
+ self.ratio = skip_ratio
+ self.CONTRAST_FACTOR = CONTRAST_FACTOR
+ self.BRIGHTNESS_FACTOR = BRIGHTNESS_FACTOR
+
+ def __call__(self, inputs):
+ return self.forward(inputs)
+
+ def forward(self, inputs):
+ skiprand = np.random.rand()
+ if skiprand < self.ratio:
+ if self.mode == 'mix':
+ # The probability of '2D' is more than '3D'
+ threshold = 1 - (1 - self.ratio) / 2
+ mode_ = '3D' if skiprand > threshold else '2D'
+ else:
+ mode_ = self.mode
+ if mode_ == '2D':
+ inputs = self.augment2D(inputs)
+ elif mode_ == '3D':
+ inputs = self.augment3D(inputs)
+ return inputs
+ else:
+ return inputs
+
+ def augment2D(self, imgs):
+ for z in range(imgs.shape[-3]):
+ img = imgs[z, :, :]
+ img *= 1 + (np.random.rand() - 0.5)*self.CONTRAST_FACTOR
+ img += (np.random.rand() - 0.5)*self.BRIGHTNESS_FACTOR
+ img = np.clip(img, 0, 1)
+ img **= 2.0**(np.random.rand()*2 - 1)
+ imgs[z, :, :] = img
+ return imgs
+
+ def augment3D(self, imgs):
+ imgs *= 1 + (np.random.rand() - 0.5)*self.CONTRAST_FACTOR
+ imgs += (np.random.rand() - 0.5)*self.BRIGHTNESS_FACTOR
+ imgs = np.clip(imgs, 0, 1)
+ imgs **= 2.0**(np.random.rand()*2 - 1)
+ return imgs
+
+
+class ElasticAugment(object):
+ '''Elasticly deform a batch. Requests larger batches upstream to avoid data
+ loss due to rotation and jitter.
+ Args:
+ control_point_spacing (``tuple`` of ``int``):
+ Distance between control points for the elastic deformation, in
+ voxels per dimension.
+ jitter_sigma (``tuple`` of ``float``):
+ Standard deviation of control point jitter distribution, in voxels
+ per dimension.
+ rotation_interval (``tuple`` of two ``floats``):
+ Interval to randomly sample rotation angles from (0, 2PI).
+ prob_slip (``float``):
+ Probability of a section to "slip", i.e., be independently moved in
+ x-y.
+ prob_shift (``float``):
+ Probability of a section and all following sections to move in x-y.
+ max_misalign (``int``):
+ Maximal voxels to shift in x and y. Samples will be drawn
+ uniformly. Used if ``prob_slip + prob_shift`` > 0.
+ subsample (``int``):
+ Instead of creating an elastic transformation on the full
+ resolution, create one subsampled by the given factor, and linearly
+ interpolate to obtain the full resolution transformation. This can
+ significantly speed up this node, at the expense of having visible
+ piecewise linear deformations for large factors. Usually, a factor
+ of 4 can savely by used without noticable changes. However, the
+ default is 1 (i.e., no subsampling).
+ '''
+ def __init__(
+ self,
+ control_point_spacing=[4, 40, 40],
+ jitter_sigma=[0,0,0], # recommend: [0, 2, 2]
+ rotation_interval=[0,0],
+ prob_slip=0, # recommend: 0.05
+ prob_shift=0, # recommend: 0.05
+ max_misalign=0, # 17 in superhuman
+ subsample=1,
+ padding=None,
+ skip_ratio=0.5): # recommend: 10
+ super(ElasticAugment, self).__init__()
+
+ self.control_point_spacing = control_point_spacing
+ self.jitter_sigma = jitter_sigma
+ self.rotation_start = rotation_interval[0]
+ self.rotation_max_amount = rotation_interval[1] - rotation_interval[0]
+ self.prob_slip = prob_slip
+ self.prob_shift = prob_shift
+ self.max_misalign = max_misalign
+ self.subsample = subsample
+ self.padding = padding
+ self.ratio = skip_ratio
+
+ def create_transformation(self, target_shape):
+
+ transformation = create_identity_transformation(
+ target_shape,
+ subsample=self.subsample)
+ # shape: channel,d,w,h
+
+ # elastic ##cost time##
+ if sum(self.jitter_sigma) > 0 and np.random.rand() < self.ratio:
+ transformation += create_elastic_transformation(
+ target_shape,
+ self.control_point_spacing,
+ self.jitter_sigma,
+ subsample=self.subsample)
+
+ # rotation = random.random()*self.rotation_max_amount + self.rotation_start
+ # if rotation != 0:
+ # transformation += create_rotation_transformation(
+ # target_shape,
+ # rotation,
+ # subsample=self.subsample)
+
+ # if self.subsample > 1:
+ # transformation = upscale_transformation(
+ # transformation,
+ # tuple(target_shape))
+
+ if self.prob_slip + self.prob_shift > 0 and np.random.rand() < self.ratio:
+ misalign(transformation, self.prob_slip,
+ self.prob_shift, self.max_misalign)
+
+ return transformation
+
+ def __call__(self, imgs, mask):
+ return self.forward(imgs, mask)
+
+ def forward(self, imgs, mask):
+ '''Args:
+ imgs: numpy array, [Z, Y, Z], it always is float and 0~1
+ mask: numpy array, [Z, Y, Z], it always is uint16
+ '''
+ if self.padding is not None:
+ imgs = np.pad(imgs, ((0,0), \
+ (self.padding,self.padding), \
+ (self.padding,self.padding)), mode='reflect')
+ mask = np.pad(mask, ((0,0), \
+ (self.padding,self.padding), \
+ (self.padding,self.padding)), mode='reflect')
+ transform = self.create_transformation(imgs.shape)
+ img_transform = apply_transformation(imgs,
+ transform,
+ interpolate=False,
+ outside_value=0, # imgs.dtype.type(-1)
+ output=np.zeros(imgs.shape, dtype=np.float32))
+ seg_transform = apply_transformation(mask,
+ transform,
+ interpolate=False,
+ outside_value=0, # mask.dtype.type(-1)
+ output=np.zeros(mask.shape, dtype=np.uint16)) # dtype=np.float32
+ # seg_transform[seg_transform < 0] = 0
+ # seg_transform[seg_transform > 60000] = 0
+ if self.padding is not None and self.padding != 0:
+ img_transform = img_transform[:, self.padding:-self.padding, self.padding:-self.padding]
+ seg_transform = seg_transform[:, self.padding:-self.padding, self.padding:-self.padding]
+ return img_transform, seg_transform
+
+
+class MissingAugment(object):
+ '''Missing section augmentation
+ Args:
+ filling: the way of filling, 'zero' or 'random'
+ mode: 'mix', 'fully' or 'partially'
+ skip_ratio: Probability of execution
+ miss_ratio: Probability of missing
+ '''
+ def __init__(self, filling='zero', mode='mix', skip_ratio=0.5, miss_ratio=0.1):
+ super(MissingAugment, self).__init__()
+ self.filling = filling
+ self.mode = mode
+ self.ratio = skip_ratio
+ self.miss_ratio = miss_ratio
+
+ def __call__(self, imgs):
+ return self.forward(imgs)
+
+ def forward(self, imgs):
+ skiprand = np.random.rand()
+ if skiprand < self.ratio:
+ if self.mode == 'mix':
+ r = np.random.rand()
+ mode_ = 'fully' if r < 0.5 else 'partially'
+ else:
+ mode_ = self.mode
+ if mode_ == 'fully':
+ imgs = self.augment_fully(imgs)
+ elif mode_ == 'partially':
+ imgs = self.augment_partially(imgs)
+ return imgs
+ else:
+ return imgs
+
+ def augment_fully(self, imgs):
+ d, h, w = imgs.shape
+ for i in range(d):
+ if np.random.rand() < self.miss_ratio:
+ if self.filling == 'zero':
+ imgs[i] = 0
+ elif self.filling == 'random':
+ imgs[i] = np.random.rand(h, w)
+ return imgs
+
+ def augment_partially(self, imgs, size_ratio=0.3):
+ d, h, w = imgs.shape
+ for i in range(d):
+ if np.random.rand() < self.miss_ratio:
+ # randomly generate an area
+ sub_h = random.randint(int(h*size_ratio), int(h*(1-size_ratio)))
+ sub_w = random.randint(int(w*size_ratio), int(w*(1-size_ratio)))
+ start_h = random.randint(0, h - sub_h - 1)
+ start_w = random.randint(0, w - sub_w - 1)
+ if self.filling == 'zero':
+ imgs[i, start_h:start_h+sub_h, start_w:start_w+sub_w] = 0
+ elif self.filling == 'random':
+ imgs[i, start_h:start_h+sub_h, start_w:start_w+sub_w] = np.random.rand(sub_h, sub_w)
+ return imgs
+
+
+class BlurAugment(object):
+ '''Out-of-focus (Blur) section augmentation
+ Args:
+ mode: 'mix', 'fully' or 'partially'
+ skip_ratio: Probability of execution
+ blur_ratio: Probability of blur
+ '''
+ def __init__(self, mode='mix', skip_ratio=0.5, blur_ratio=0.1):
+ super(BlurAugment, self).__init__()
+ self.mode = mode
+ self.ratio = skip_ratio
+ self.blur_ratio = blur_ratio
+
+ def __call__(self, imgs):
+ return self.forward(imgs)
+
+ def forward(self, imgs):
+ skiprand = np.random.rand()
+ if skiprand < self.ratio:
+ if self.mode == 'mix':
+ r = np.random.rand()
+ mode_ = 'fully' if r < 0.5 else 'partially'
+ else:
+ mode_ = self.mode
+ if mode_ == 'fully':
+ imgs = self.augment_fully(imgs)
+ elif mode_ == 'partially':
+ imgs = self.augment_partially(imgs)
+ return imgs
+ else:
+ return imgs
+
+ def augment_fully(self, imgs):
+ d, h, w = imgs.shape
+ for i in range(d):
+ if np.random.rand() < self.blur_ratio:
+ sigma = np.random.uniform(0, 5)
+ imgs[i] = gaussian_filter(imgs[i], sigma)
+ return imgs
+
+ def augment_partially(self, imgs, size_ratio=0.3):
+ d, h, w = imgs.shape
+ for i in range(d):
+ if np.random.rand() < self.blur_ratio:
+ # randomly generate an area
+ sub_h = random.randint(int(h*size_ratio), int(h*(1-size_ratio)))
+ sub_w = random.randint(int(w*size_ratio), int(w*(1-size_ratio)))
+ start_h = random.randint(0, h - sub_h - 1)
+ start_w = random.randint(0, w - sub_w - 1)
+ sigma = np.random.uniform(0, 5)
+ imgs[i, start_h:start_h+sub_h, start_w:start_w+sub_w] = \
+ gaussian_filter(imgs[i, start_h:start_h+sub_h, start_w:start_w+sub_w], sigma)
+ return imgs
+
+
+def show(img3d):
+ # only used for image with shape [18, 160, 160]
+ row = 4
+ column = 5
+ num = 18
+ size = 160
+ img_all = np.zeros((size*row, size*column), dtype=np.uint8)
+ for i in range(row):
+ for j in range(column):
+ index = i*column + j
+ if index >= num:
+ img = np.zeros_like(img3d[0], dtype=np.uint8)
+ else:
+ img = (img3d[index] * 255).astype(np.uint8)
+ img_all[i*size:(i+1)*size, j*size:(j+1)*size] = img
+ return img_all
+
+
+def show_lb(img3d):
+ # only used for image with shape [18, 160, 160]
+ row = 4
+ column = 5
+ num = 18
+ size = 160
+ ids = np.unique(img3d)
+ color_pred = np.zeros([num, size, size, 3], dtype=np.uint8)
+ idx = np.searchsorted(ids, img3d)
+ for i in range(3):
+ color_val = np.random.randint(0, 255, ids.shape)
+ if ids[0] == 0:
+ color_val[0] = 0
+ color_pred[:,:,:,i] = color_val[idx]
+
+ img_all = np.zeros((size*row, size*column, 3), dtype=np.uint8)
+ for i in range(row):
+ for j in range(column):
+ index = i*column + j
+ if index >= num:
+ img = np.zeros_like((size, size, 3), dtype=np.uint8)
+ else:
+ img = color_pred[index]
+ img_all[i*size:(i+1)*size, j*size:(j+1)*size, :] = img
+ return img_all
+
+
+def elastic_deform_3d_cuda(image_in,
+ label_in,
+ prob,
+ random_state=None,
+ padding=20,
+ alpha=(10,50),
+ sigma=10,
+ device='cuda:0'):
+ """Elastic deformation of image_ins as described in [Simard2003]_.
+ .. [Simard2003] Simard, Steinkraus and Platt, "Best Practices for
+ Convolutional Neural Networks applied to Visual Document Analysis", in
+ Proc. of the International Conference on Document Analysis and
+ Recognition, 2003.
+ """
+ # skip
+ if random.uniform(0, 1) > prob:
+ return image_in, label_in
+
+ if padding is not None:
+ image_in = np.pad(image_in, ((0,0), \
+ (padding, padding), \
+ (padding, padding)), mode='reflect')
+ label_in = np.pad(label_in, ((0,0), \
+ (padding, padding), \
+ (padding, padding)), mode='reflect')
+
+ alpha = np.random.uniform(alpha[0], alpha[1])
+ if random_state is None:
+ random_state = np.random.RandomState(None)
+
+ shape = image_in.shape
+
+ #rdx = torch.Tensor(random_state.rand(*shape) * 2 - 1).unsqueeze(0).unsqueeze(0).to(self.device)
+ #rdy = torch.Tensor(random_state.rand(*shape) * 2 - 1).unsqueeze(0).unsqueeze(0).to(self.device)
+ #rdz = torch.Tensor(random_state.rand(*shape) * 2 - 1).unsqueeze(0).unsqueeze(0).to(self.device)
+ #dx = self.gaussian_filter(rdx) * alpha
+ #dy = self.gaussian_filter(rdy) * alpha
+ #dz = self.gaussian_filter(rdz) * alpha
+
+ #dx = np.squeeze(dx.data.cpu().numpy())
+ #dy = np.squeeze(dy.data.cpu().numpy())
+ #dz = np.squeeze(dz.data.cpu().numpy())
+
+ dx = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, order=0, mode='constant', cval=0) * alpha
+ dy = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, order=0, mode='constant', cval=0) * alpha
+ dz = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, order=0, mode='constant', cval=0) * alpha
+
+ grid_x, grid_y, grid_z = np.meshgrid(np.arange(shape[1]), np.arange(shape[0]), np.arange(shape[2]))
+ def_grid_x = grid_x + dx
+ def_grid_y = grid_y + dy
+ def_grid_z = grid_z + dz
+
+ gx = 2.0*def_grid_x/(shape[2]-1)-1.
+ gy = 2.0*def_grid_y/(shape[1]-1)-1.
+ gz = 2.0*def_grid_z/(shape[0]-1)-1.
+
+ #indices = np.reshape(def_grid_y, (-1, 1)), np.reshape(def_grid_x, (-1, 1)), np.reshape(def_grid_z, (-1, 1))
+ #out = map_coordinates(image_in, indices, order=1).reshape(shape)
+
+ # torch_grid = torch.Tensor(np.stack((gx,gy,gz),3)).unsqueeze(0).to(device)
+ torch_grid = torch.Tensor(np.stack((gz,gx,gy),3)).unsqueeze(0).to(device)
+ torch_im = torch.Tensor(np.expand_dims(np.expand_dims(image_in, axis=0), axis=0).copy()).to(device)
+ torch_lb = torch.Tensor(np.expand_dims(label_in, axis=0).copy()).to(device)
+ with torch.no_grad():
+ torch_im_out = torch.nn.functional.grid_sample(torch_im, torch_grid, mode='bilinear', padding_mode='zeros')
+ torch_lb_out = torch.nn.functional.grid_sample(torch_lb, torch_grid, mode='bilinear', padding_mode='zeros')
+
+ image_out = np.squeeze(torch_im_out.data.cpu().numpy()).astype(np.uint8)
+ label_out = np.squeeze(torch_lb_out.data.cpu().numpy()).astype(np.uint8)
+
+ if padding is not None and padding != 0:
+ image_out = image_out[:, padding:-padding, padding:-padding]
+ label_out = label_out[:, padding:-padding, padding:-padding]
+ return image_out, label_out
+
+if __name__ == "__main__":
+ import os
+ import cv2
+ import h5py
+
+ input_vol = '../data/snemi3d/train-input.h5'
+ f = h5py.File(input_vol, 'r')
+ raw = f['main'][:]
+ f.close()
+
+ input_vol = '../data/snemi3d/train-labels.h5'
+ f = h5py.File(input_vol, 'r')
+ lbs = f['main'][:]
+ f.close()
+
+ out = './debug_img'
+ raw = raw.astype(np.float32) / 255.0
+ vol = raw[0:18, 0:160, 0:160]
+ lb = lbs[0:18, 0:160, 0:160]
+ # vol_img = show_lb(lb)
+ # cv2.imwrite(os.path.join(out, 'raw.png'), vol_img)
+
+ ##################################################
+ # Data_aug = ElasticAugment(jitter_sigma=[0,2,2],
+ # prob_slip=0.5,
+ # prob_shift=0.5,
+ # max_misalign=17,
+ # padding=20)
+ # print('min=%d, max=%d' % (np.min(lb), np.max(lb)))
+ ##################################################
+ # Data_aug = MissingAugment(filling='random')
+ ##################################################
+ Data_aug = BlurAugment(blur_ratio=0.1)
+ for i in range(20):
+ # vol_aug, lb_aug = Data_aug(vol.copy(), lb.copy())
+ # print('min=%d, max=%d' % (np.min(lb_aug), np.max(lb_aug)))
+ # print(lb_aug.dtype)
+ # vol_img = show_lb(lb_aug)
+ vol_aug = Data_aug(vol.copy())
+ vol_img = show(vol_aug)
+
+ cv2.imwrite(os.path.join(out, 'raw_aug'+str(i)+'.png'), vol_img)
+ print('Done')
\ No newline at end of file
diff --git a/legacy/Train_and_Inference/utils/augmentation_affine.py b/legacy/Train_and_Inference/utils/augmentation_affine.py
new file mode 100644
index 0000000..f863b63
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/augmentation_affine.py
@@ -0,0 +1,240 @@
+import cv2
+import math
+import numpy as np
+
+from utils import affine
+
+class SegCVTransformRandomCropRotateScale(object):
+ """
+ Random crop with random scale.
+ """
+ def __init__(self, crop_size, crop_offset, rot_mag, max_scale, uniform_scale=True, constrain_rot_scale=True,
+ rng=None):
+ if crop_offset is None:
+ crop_offset = [0, 0]
+ self.crop_size = tuple(crop_size)
+ self.crop_size_arr = np.array(crop_size)
+ self.crop_offset = np.array(crop_offset)
+ self.rot_mag_rad = math.radians(rot_mag)
+ self.log_max_scale = np.log(max_scale)
+ self.uniform_scale = uniform_scale
+ self.constrain_rot_scale = constrain_rot_scale
+ self.__rng = rng
+
+ @property
+ def rng(self):
+ if self.__rng is None:
+ self.__rng = np.random.RandomState()
+ return self.__rng
+
+ def transform_single(self, sample0):
+ sample0 = sample0.copy()
+
+ # Extract contents
+ image = sample0['image_arr']
+
+ # Choose scale and rotation
+ if self.uniform_scale:
+ scale_factor_yx = np.exp(self.rng.uniform(-self.log_max_scale, self.log_max_scale, size=(1,)))
+ scale_factor_yx = np.repeat(scale_factor_yx, 2, axis=0)
+ else:
+ scale_factor_yx = np.exp(self.rng.uniform(-self.log_max_scale, self.log_max_scale, size=(2,)))
+ rot_theta = self.rng.uniform(-self.rot_mag_rad, self.rot_mag_rad, size=(1,))
+
+ # Scale the crop size by the inverse of the scale
+ sc_size = self.crop_size_arr / scale_factor_yx
+
+ # Randomly choose centre
+ img_size = np.array(image.shape[:2])
+ extra = np.maximum(img_size - sc_size, 0.0)
+ centre = extra * self.rng.uniform(0.0, 1.0, size=(2,)) + np.minimum(sc_size, img_size) * 0.5
+
+ # Build affine transformation matrix
+ local_xf = affine.cat_nx2x3(
+ affine.translation_matrices(self.crop_size_arr[None, ::-1] * 0.5),
+ affine.rotation_matrices(rot_theta),
+ affine.scale_matrices(scale_factor_yx[None, ::-1]),
+ affine.translation_matrices(-centre[None, ::-1]),
+ )
+
+ # Reflect the image
+ # Use nearest neighbour sampling to stay consistent with labels, if labels present
+ if 'labels_arr' in sample0:
+ interpolation = cv2.INTER_NEAREST
+ else:
+ interpolation = self.rng.choice([cv2.INTER_NEAREST, cv2.INTER_LINEAR])
+
+ sample0['image_arr'] = cv2.warpAffine(image, local_xf[0], self.crop_size[::-1], flags=interpolation, borderValue=0, borderMode=cv2.BORDER_REFLECT_101)
+
+ # Don't reflect labels and mask
+ if 'labels_arr' in sample0:
+ sample0['labels_arr'] = cv2.warpAffine(sample0['labels_arr'], local_xf[0], self.crop_size[::-1], flags=cv2.INTER_NEAREST, borderValue=0, borderMode=cv2.BORDER_CONSTANT)
+
+ if 'mask_arr' in sample0:
+ sample0['mask_arr'] = cv2.warpAffine(sample0['mask_arr'], local_xf[0], self.crop_size[::-1], flags=interpolation, borderValue=0, borderMode=cv2.BORDER_CONSTANT)
+
+ if 'xf_cv' in sample0:
+ sample0['xf_cv'] = affine.cat_nx2x3(local_xf, sample0['xf_cv'][None, ...])[0]
+
+ return sample0
+
+ def transform_pair(self, sample0, sample1):
+ sample0 = sample0.copy()
+ sample1 = sample1.copy()
+
+ # Choose scales and rotations
+ if self.constrain_rot_scale:
+ if self.uniform_scale:
+ scale_factors_yx = np.exp(self.rng.uniform(-self.log_max_scale, self.log_max_scale, size=(1, 1)))
+ scale_factors_yx = np.repeat(scale_factors_yx, 2, axis=1)
+ else:
+ scale_factors_yx = np.exp(self.rng.uniform(-self.log_max_scale, self.log_max_scale, size=(1, 2)))
+
+ rot_thetas = self.rng.uniform(-self.rot_mag_rad, self.rot_mag_rad, size=(1,))
+ scale_factors_yx = np.repeat(scale_factors_yx, 2, axis=0)
+ rot_thetas = np.repeat(rot_thetas, 2, axis=0)
+ else:
+ if self.uniform_scale:
+ scale_factors_yx = np.exp(self.rng.uniform(-self.log_max_scale, self.log_max_scale, size=(2, 1)))
+ scale_factors_yx = np.repeat(scale_factors_yx, 2, axis=1)
+ else:
+ scale_factors_yx = np.exp(self.rng.uniform(-self.log_max_scale, self.log_max_scale, size=(2, 2)))
+ rot_thetas = self.rng.uniform(-self.rot_mag_rad, self.rot_mag_rad, size=(2,))
+
+ img_size = np.array(sample0['image_arr'].shape[:2])
+
+ # Scale the crop size by the inverse of the scale
+ sc_size = self.crop_size_arr / scale_factors_yx.min(axis=0)
+ crop_centre_pos = np.minimum(sc_size, img_size) * 0.5
+
+ # Randomly choose centres
+ extra = np.maximum(img_size - sc_size, 0.0)
+ centre0 = extra * self.rng.uniform(0.0, 1.0, size=(2,)) + crop_centre_pos
+ offset1 = np.round(self.crop_offset * self.rng.uniform(-1.0, 1.0, size=(2,)))
+ centre_xlat = np.stack([centre0, centre0], axis=0)
+ offset1_xlat = np.stack([np.zeros((2,)), offset1], axis=0)
+
+ # Build affine transformation matrices
+ local_xfs = affine.cat_nx2x3(
+ affine.translation_matrices(self.crop_size_arr[None, ::-1] * 0.5),
+ affine.translation_matrices(offset1_xlat[:, ::-1]),
+ affine.rotation_matrices(rot_thetas),
+ affine.scale_matrices(scale_factors_yx[:, ::-1]),
+ affine.translation_matrices(-centre_xlat[:, ::-1]),
+ )
+
+ # Use nearest neighbour sampling to stay consistent with labels, if labels present
+ interpolation = cv2.INTER_NEAREST if 'labels_arr' in sample0 else cv2.INTER_LINEAR
+ sample0['image_arr'] = cv2.warpAffine(sample0['image_arr'], local_xfs[0], self.crop_size[::-1], flags=interpolation,
+ borderValue=0, borderMode=cv2.BORDER_REFLECT_101)
+ sample1['image_arr'] = cv2.warpAffine(sample1['image_arr'], local_xfs[1], self.crop_size[::-1], flags=interpolation,
+ borderValue=0, borderMode=cv2.BORDER_REFLECT_101)
+
+ if 'labels_arr' in sample0:
+ sample0['labels_arr'] = cv2.warpAffine(sample0['labels_arr'], local_xfs[0], self.crop_size[::-1], flags=cv2.INTER_NEAREST,
+ borderValue=0, borderMode=cv2.BORDER_CONSTANT)
+ sample1['labels_arr'] = cv2.warpAffine(sample1['labels_arr'], local_xfs[1], self.crop_size[::-1], flags=cv2.INTER_NEAREST,
+ borderValue=0, borderMode=cv2.BORDER_CONSTANT)
+
+ if 'mask_arr' in sample0:
+ sample0['mask_arr'] = cv2.warpAffine(sample0['mask_arr'], local_xfs[0], self.crop_size[::-1], flags=interpolation,
+ borderValue=0, borderMode=cv2.BORDER_CONSTANT)
+ sample1['mask_arr'] = cv2.warpAffine(sample1['mask_arr'], local_xfs[1], self.crop_size[::-1], flags=interpolation,
+ borderValue=0, borderMode=cv2.BORDER_CONSTANT)
+
+ if 'xf_cv' in sample0:
+ xf01 = affine.cat_nx2x3(local_xfs, np.stack([sample0['xf_cv'], sample1['xf_cv']], axis=0))
+ sample0['xf_cv'] = xf01[0]
+ sample1['xf_cv'] = xf01[1]
+
+ return sample0, sample1
+
+
+class SegCVTransformRandomFlip(object):
+ def __init__(self, hflip, vflip, hvflip, rng=None):
+ self.hflip = hflip
+ self.vflip = vflip
+ self.hvflip = hvflip
+ self.__rng = rng
+
+ @property
+ def rng(self):
+ if self.__rng is None:
+ self.__rng = np.random.RandomState()
+ return self.__rng
+
+ @staticmethod
+ def flip_image(img, flip_xyd):
+ if flip_xyd[0]:
+ img = img[:, ::-1]
+ if flip_xyd[1]:
+ img = img[::-1, ...]
+ if flip_xyd[2]:
+ img = np.swapaxes(img, 0, 1)
+ return img.copy()
+
+ def transform_single(self, sample):
+ sample = sample.copy()
+
+ # Flip flags
+ flip_flags_xyd = self.rng.binomial(1, 0.5, size=(3,)) != 0
+ flip_flags_xyd = flip_flags_xyd & np.array([self.hflip, self.vflip, self.hvflip])
+
+ sample['image_arr'] = self.flip_image(sample['image_arr'], flip_flags_xyd)
+
+ if 'mask_arr' in sample:
+ sample['mask_arr'] = self.flip_image(sample['mask_arr'], flip_flags_xyd)
+
+ if 'labels_arr' in sample:
+ sample['labels_arr'] = self.flip_image(sample['labels_arr'], flip_flags_xyd)
+
+ if 'xf_cv' in sample:
+ sample['xf_cv'] = affine.cat_nx2x3(
+ affine.flip_xyd_matrices(flip_flags_xyd[None, ...], sample['image_arr'].shape[:2]),
+ sample['xf_cv'][None, ...],
+ )[0]
+
+ return sample
+
+ def transform_pair(self, sample0, sample1):
+ sample0 = sample0.copy()
+ sample1 = sample1.copy()
+
+ # Flip flags
+ flip_flags_xyd = self.rng.binomial(1, 0.5, size=(2, 3)) != 0
+ flip_flags_xyd = flip_flags_xyd & np.array([[self.hflip, self.vflip, self.hvflip]])
+
+ sample0['image_arr'] = self.flip_image(sample0['image_arr'], flip_flags_xyd[0])
+ sample1['image_arr'] = self.flip_image(sample1['image_arr'], flip_flags_xyd[1])
+
+ if 'mask_arr' in sample0:
+ sample0['mask_arr'] = self.flip_image(sample0['mask_arr'], flip_flags_xyd[0])
+ sample1['mask_arr'] = self.flip_image(sample1['mask_arr'], flip_flags_xyd[1])
+
+ if 'labels_arr' in sample0:
+ sample0['labels_arr'] = self.flip_image(sample0['labels_arr'], flip_flags_xyd[0])
+ sample1['labels_arr'] = self.flip_image(sample1['labels_arr'], flip_flags_xyd[1])
+
+ if 'xf_cv' in sample0:
+ # False -> 1, True -> -1
+ flip_scale_xy = flip_flags_xyd[:, :2] * -2 + 1
+ # Negative scale factors need to be combined with a translation whose value is (image_size - 1)
+ # Mask the translation with the flip flags to only apply it where flipping is done
+ flip_xlat_xy = flip_flags_xyd[:, :2] * (np.array([sample0['image_arr'].shape[:2][::-1],
+ sample1['image_arr'].shape[:2][::-1]]).astype(float) - 1)
+
+ hv_flip_xf = affine.identity_xf(2)
+ hv_flip_xf[flip_flags_xyd[:, 2]] = hv_flip_xf[flip_flags_xyd[:, 2], ::-1, :]
+
+ xf01 = np.stack([sample0['xf_cv'], sample1['xf_cv']], axis=0)
+ xf01 = affine.cat_nx2x3(
+ hv_flip_xf,
+ affine.translation_matrices(flip_xlat_xy),
+ affine.scale_matrices(flip_scale_xy),
+ xf01,
+ )
+ sample0['xf_cv'] = xf01[0]
+ sample1['xf_cv'] = xf01[1]
+
+ return sample0, sample1
+
diff --git a/legacy/Train_and_Inference/utils/compute_sdf.py b/legacy/Train_and_Inference/utils/compute_sdf.py
new file mode 100644
index 0000000..a8a1505
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/compute_sdf.py
@@ -0,0 +1,32 @@
+import numpy as np
+from scipy.ndimage import distance_transform_edt as distance
+from skimage import segmentation as skimage_seg
+
+def compute_sdf(img_gt):
+ """
+ compute the signed distance map of binary mask
+ input: segmentation, shape = (batch_size, x, y, z)
+ output: the Signed Distance Map (SDM)
+ sdf(x) = 0; x in segmentation boundary
+ -inf|x-y|; x in segmentation
+ +inf|x-y|; x out of segmentation
+ normalize sdf to [-1,1]
+ """
+ normalized_sdf = np.zeros_like(img_gt, dtype=np.float32)
+ ids, counts = np.unique(img_gt, return_counts=True)
+ # remove id 0
+ if ids[0] == 0:
+ ids = ids[1:]
+ # if ids is None
+ if len(ids) == 0:
+ return normalized_sdf
+
+ for id in ids:
+ posmask = np.zeros_like(img_gt)
+ posmask[img_gt == id] = 1
+ posmask = posmask.astype(np.bool)
+ if posmask.any():
+ posdis = distance(posmask)
+ posdis = (posdis - posdis.min()) / (posdis.max() - posdis.min())
+ normalized_sdf += posdis
+ return normalized_sdf
diff --git a/legacy/Train_and_Inference/utils/consistency_aug.py b/legacy/Train_and_Inference/utils/consistency_aug.py
new file mode 100644
index 0000000..9ee32db
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/consistency_aug.py
@@ -0,0 +1,235 @@
+import cv2
+import torch
+import random
+import numpy as np
+import torch.nn.functional as F
+
+def simple_augment(data, rule):
+ assert np.size(rule) == 4
+ assert data.ndim == 3
+ # z reflection
+ if rule[0]:
+ data = data[::-1, :, :]
+ # x reflection
+ if rule[1]:
+ data = data[:, :, ::-1]
+ # y reflection
+ if rule[2]:
+ data = data[:, ::-1, :]
+ # transpose in xy
+ if rule[3]:
+ data = np.transpose(data, (0, 2, 1))
+ return data
+
+def simple_augment_torch(data, rule):
+ assert np.size(rule) == 4
+ assert len(data.shape) == 4
+ # z reflection
+ if rule[0]:
+ data = torch.flip(data, [1])
+ # x reflection
+ if rule[1]:
+ data = torch.flip(data, [3])
+ # y reflection
+ if rule[2]:
+ data = torch.flip(data, [2])
+ # transpose in xy
+ if rule[3]:
+ data = data.permute(0, 1, 3, 2)
+ return data
+
+def simple_augment_reverse(data, rule):
+ assert np.size(rule) == 4
+ assert len(data.shape) == 5
+ # transpose in xy
+ if rule[3]:
+ # data = np.transpose(data, (0, 1, 2, 4, 3))
+ data = data.permute(0, 1, 2, 4, 3)
+ # y reflection
+ if rule[2]:
+ # data = data[:, :, :, ::-1, :]
+ data = torch.flip(data, [3])
+ # x reflection
+ if rule[1]:
+ # data = data[:, :, :, :, ::-1]
+ data = torch.flip(data, [4])
+ # z reflection
+ if rule[0]:
+ # data = data[:, :, ::-1, :, :]
+ data = torch.flip(data, [2])
+ return data
+
+def order_aug(imgs, num_patch=4):
+ assert imgs.shape[-1] % num_patch == 0
+ patch_size = imgs.shape[-1] // num_patch
+ new_imgs = np.zeros_like(imgs, dtype=np.float32)
+ # ran_order = np.random.shuffle(np.arange(num_patch**2))
+ ran_order = np.random.permutation(num_patch**2)
+ for k in range(num_patch**2):
+ xid_new = k // num_patch
+ yid_new = k % num_patch
+ order_id = ran_order[k]
+ xid_old = order_id // num_patch
+ yid_old = order_id % num_patch
+ new_imgs[:, xid_new*patch_size:(xid_new+1)*patch_size, yid_new*patch_size:(yid_new+1)*patch_size] = \
+ imgs[:, xid_old*patch_size:(xid_old+1)*patch_size, yid_old*patch_size:(yid_old+1)*patch_size]
+ return new_imgs
+
+def gen_mask(imgs, model_type='superhuman', min_mask_counts=40, max_mask_counts=60, min_mask_size=[3, 5, 5], max_mask_size=[7, 20, 20]):
+ if model_type == 'mala':
+ net_crop_size = [14, 106, 106]
+ else:
+ net_crop_size = [0, 0, 0]
+ crop_size = list(imgs.shape)
+ mask = np.ones_like(imgs, dtype=np.float32)
+ mask_counts = random.randint(min_mask_counts, max_mask_counts)
+ mask_size_z = random.randint(min_mask_size[0], max_mask_size[0])
+ mask_size_xy = random.randint(min_mask_size[1], max_mask_size[1])
+ for k in range(mask_counts):
+ mz = random.randint(net_crop_size[0], crop_size[0]-mask_size_z-net_crop_size[0])
+ my = random.randint(net_crop_size[1], crop_size[1]-mask_size_xy-net_crop_size[1])
+ mx = random.randint(net_crop_size[2], crop_size[2]-mask_size_xy-net_crop_size[2])
+ mask[mz:mz+mask_size_z, my:my+mask_size_xy, mx:mx+mask_size_xy] = 0
+ return mask
+
+def resize_3d(imgs, det_size, mode='linear'):
+ new_imgs = []
+ for k in range(imgs.shape[0]):
+ temp = imgs[k]
+ if mode == 'linear':
+ temp = cv2.resize(temp, (det_size, det_size), interpolation=cv2.INTER_LINEAR)
+ elif mode == 'nearest':
+ temp = cv2.resize(temp, (det_size, det_size), interpolation=cv2.INTER_NEAREST)
+ else:
+ raise AttributeError('No this interpolation mode!')
+ new_imgs.append(temp)
+ new_imgs = np.asarray(new_imgs)
+ return new_imgs
+
+def add_gauss_noise(imgs, min_std=0.1, max_std=0.1, norm_mode='norm'):
+ if min_std == max_std:
+ std = min_std
+ else:
+ std = random.uniform(min_std, max_std)
+ gaussian = np.random.normal(0, std, (imgs.shape))
+ imgs = imgs + gaussian
+ if norm_mode == 'norm':
+ imgs = (imgs-np.min(imgs)) / (np.max(imgs)-np.min(imgs))
+ elif norm_mode == 'trunc':
+ imgs[imgs<0] = 0
+ imgs[imgs>1] = 1
+ else:
+ pass
+ return imgs
+
+def add_gauss_blur(imgs, kernel_size=5, sigma=0):
+ outs = []
+ for k in range(imgs.shape[0]):
+ temp = imgs[k]
+ temp = cv2.GaussianBlur(temp, (kernel_size,kernel_size), sigma)
+ outs.append(temp)
+ outs = np.asarray(outs, dtype=np.float32)
+ outs[outs < 0] = 0
+ outs[outs > 1] = 1
+ return outs
+
+def add_intensity(imgs, contrast_factor=0.1, brightness_factor=0.1):
+ # imgs *= 1 + (np.random.rand() - 0.5) * contrast_factor
+ # imgs += (np.random.rand() - 0.5) * brightness_factor
+ # imgs = np.clip(imgs, 0, 1)
+ # imgs **= 2.0**(np.random.rand()*2 - 1)
+ imgs *= 1 + contrast_factor
+ imgs += brightness_factor
+ imgs = np.clip(imgs, 0, 1)
+ return imgs
+
+def interp_5d(data, det_size, mode='bilinear'):
+ assert len(data.shape) == 5, "the dimension of data must be 5!"
+ out = []
+ depth = data.shape[2]
+ for k in range(depth):
+ temp = data[:,:,k,:,:]
+ if mode == 'bilinear':
+ temp = F.interpolate(temp, size=(det_size, det_size), mode='bilinear', align_corners=True)
+ elif mode == 'nearest':
+ temp = F.interpolate(temp, size=(det_size, det_size), mode='nearest')
+ out.append(temp)
+ out = torch.stack(out, dim=2)
+ return out
+
+def convert_consistency_scale(gt, det_size):
+ B, C, D, H, W = gt.shape
+ gt = gt.detach().clone()
+ out_gt = []
+ masks = []
+ for k in range(B):
+ gt_temp = gt[k]
+ det_size_temp = det_size[k]
+ if det_size_temp[0] == gt_temp.shape[-1]:
+ mask = torch.ones_like(gt_temp)
+ out_gt.append(gt_temp)
+ masks.append(mask)
+ elif det_size_temp[0] > gt_temp.shape[-1]:
+ shift = int((det_size_temp[0] - gt_temp.shape[-1]) // 2)
+ gt_padding = torch.zeros((1, C, D, int(det_size_temp[0]), int(det_size_temp[0]))).float().cuda()
+ mask = torch.zeros_like(gt_padding)
+ gt_padding[0,:,:,shift:-shift,shift:-shift] = gt_temp
+ mask[0,:,:,shift:-shift,shift:-shift] = 1
+ # gt_padding = F.interpolate(gt_padding, size=(D, int(gt_temp.shape[-1]), int(gt_temp.shape[-1])), mode='trilinear', align_corners=True)
+ gt_padding = interp_5d(gt_padding, int(gt_temp.shape[-1]), mode='bilinear')
+ mask = F.interpolate(mask, size=(D, int(gt_temp.shape[-1]), int(gt_temp.shape[-1])), mode='nearest')
+ gt_padding = torch.squeeze(gt_padding, dim=0)
+ mask = torch.squeeze(mask, dim=0)
+ out_gt.append(gt_padding)
+ masks.append(mask)
+ else:
+ shift = int((gt_temp.shape[-1] - det_size_temp[0]) // 2)
+ mask = torch.zeros_like(gt_temp)
+ mask[:,:,shift:-shift,shift:-shift] = 1
+ gt_padding = gt_temp[:,:,shift:-shift,shift:-shift]
+ gt_padding = gt_padding[None, ...]
+ # gt_padding = F.interpolate(gt_padding, size=(D, int(gt_temp.shape[-1]), int(gt_temp.shape[-1])), mode='trilinear', align_corners=True)
+ gt_padding = interp_5d(gt_padding, int(gt_temp.shape[-1]), mode='bilinear')
+ gt_padding = torch.squeeze(gt_padding, dim=0)
+ out_gt.append(gt_padding)
+ masks.append(mask)
+ out_gt = torch.stack(out_gt, dim=0)
+ masks = torch.stack(masks, dim=0)
+ return out_gt, masks
+
+def convert_consistency_flip(gt, rules):
+ B, C, D, H, W = gt.shape
+ gt = gt.detach().clone()
+ rules = rules.data.cpu().numpy()
+ out_gt = []
+ for k in range(B):
+ gt_temp = gt[k]
+ rule = rules[k]
+ gt_temp = simple_augment_torch(gt_temp, rule)
+ out_gt.append(gt_temp)
+ out_gt = torch.stack(out_gt, dim=0)
+ return out_gt
+
+def convert_consistency_rot(gt, rotnums):
+ B, C, D, H, W = gt.shape
+ gt = gt.detach().clone()
+ rotnums = rotnums.data.cpu().numpy()
+ out_gt = []
+ for k in range(B):
+ gt_temp = gt[k]
+ rotnum = int(rotnums[k])
+ gt_temp = torch.rot90(gt_temp, rotnum, [2,3])
+ out_gt.append(gt_temp)
+ out_gt = torch.stack(out_gt, dim=0)
+ return out_gt
+
+
+if __name__ == "__main__":
+ test = np.random.random((3,3,18,160,160)).astype(np.float32)
+ det_size = np.asarray([[160],[320],[80]], dtype=np.float32)
+ test = torch.tensor(test).to('cuda:0')
+ det_size = torch.tensor(det_size).to('cuda:0')
+ out_gt, masks = convert_consistency_scale(test, det_size)
+ print(out_gt.shape)
+
+
diff --git a/legacy/Train_and_Inference/utils/consistency_aug_perturbations.py b/legacy/Train_and_Inference/utils/consistency_aug_perturbations.py
new file mode 100644
index 0000000..e1a2d6d
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/consistency_aug_perturbations.py
@@ -0,0 +1,728 @@
+import cv2
+import torch
+import random
+import numpy as np
+import torch.nn.functional as F
+from skimage import filters
+from scipy.ndimage.filters import gaussian_filter
+
+from utils.augmentation import create_identity_transformation
+from utils.augmentation import create_elastic_transformation
+from utils.augmentation import apply_transformation
+from utils.augmentation import misalign
+from utils.flow_synthesis import gen_line, gen_flow
+from utils.image_warp import image_warp
+
+
+def simple_augment(data, rule):
+ assert np.size(rule) == 4
+ assert data.ndim == 3
+ # z reflection
+ if rule[0]:
+ data = data[::-1, :, :]
+ # x reflection
+ if rule[1]:
+ data = data[:, :, ::-1]
+ # y reflection
+ if rule[2]:
+ data = data[:, ::-1, :]
+ # transpose in xy
+ if rule[3]:
+ data = np.transpose(data, (0, 2, 1))
+ return data
+
+
+def simple_augment_torch(data, rule):
+ assert np.size(rule) == 4
+ assert len(data.shape) == 4
+ # z reflection
+ if rule[0]:
+ data = torch.flip(data, [1])
+ # x reflection
+ if rule[1]:
+ data = torch.flip(data, [3])
+ # y reflection
+ if rule[2]:
+ data = torch.flip(data, [2])
+ # transpose in xy
+ if rule[3]:
+ data = data.permute(0, 1, 3, 2)
+ return data
+
+
+def simple_augment_reverse(data, rule):
+ assert np.size(rule) == 4
+ assert len(data.shape) == 5
+ # transpose in xy
+ if rule[3]:
+ # data = np.transpose(data, (0, 1, 2, 4, 3))
+ data = data.permute(0, 1, 2, 4, 3)
+ # y reflection
+ if rule[2]:
+ # data = data[:, :, :, ::-1, :]
+ data = torch.flip(data, [3])
+ # x reflection
+ if rule[1]:
+ # data = data[:, :, :, :, ::-1]
+ data = torch.flip(data, [4])
+ # z reflection
+ if rule[0]:
+ # data = data[:, :, ::-1, :, :]
+ data = torch.flip(data, [2])
+ return data
+
+
+def order_aug(imgs, num_patch=4):
+ assert imgs.shape[-1] % num_patch == 0
+ patch_size = imgs.shape[-1] // num_patch
+ new_imgs = np.zeros_like(imgs, dtype=np.float32)
+ # ran_order = np.random.shuffle(np.arange(num_patch**2))
+ ran_order = np.random.permutation(num_patch ** 2)
+ for k in range(num_patch ** 2):
+ xid_new = k // num_patch
+ yid_new = k % num_patch
+ order_id = ran_order[k]
+ xid_old = order_id // num_patch
+ yid_old = order_id % num_patch
+ new_imgs[:, xid_new * patch_size:(xid_new + 1) * patch_size, yid_new * patch_size:(yid_new + 1) * patch_size] = \
+ imgs[:, xid_old * patch_size:(xid_old + 1) * patch_size, yid_old * patch_size:(yid_old + 1) * patch_size]
+ return new_imgs
+
+
+def gen_mask(imgs, net_crop_size=[0, 0, 0], mask_counts=80, mask_size_z=8, mask_size_xy=15):
+ crop_size = list(imgs.shape)
+ mask = np.ones_like(imgs, dtype=np.float32)
+ for k in range(mask_counts):
+ mz = random.randint(net_crop_size[0], crop_size[0] - mask_size_z - net_crop_size[0])
+ my = random.randint(net_crop_size[1], crop_size[1] - mask_size_xy - net_crop_size[1])
+ mx = random.randint(net_crop_size[2], crop_size[2] - mask_size_xy - net_crop_size[2])
+ mask[mz:mz + mask_size_z, my:my + mask_size_xy, mx:mx + mask_size_xy] = 0
+ return mask
+
+
+def resize_3d(imgs, det_size, mode='linear'):
+ new_imgs = []
+ for k in range(imgs.shape[0]):
+ temp = imgs[k]
+ if mode == 'linear':
+ temp = cv2.resize(temp, (det_size, det_size), interpolation=cv2.INTER_LINEAR)
+ elif mode == 'nearest':
+ temp = cv2.resize(temp, (det_size, det_size), interpolation=cv2.INTER_NEAREST)
+ else:
+ raise AttributeError('No this interpolation mode!')
+ new_imgs.append(temp)
+ new_imgs = np.asarray(new_imgs)
+ return new_imgs
+
+
+def add_gauss_noise(imgs, std=0.01, norm_mode='norm'):
+ gaussian = np.random.normal(0, std, (imgs.shape))
+ imgs = imgs + gaussian
+ if norm_mode == 'norm':
+ imgs = (imgs - np.min(imgs)) / (np.max(imgs) - np.min(imgs))
+ elif norm_mode == 'trunc':
+ imgs[imgs < 0] = 0
+ imgs[imgs > 1] = 1
+ else:
+ raise NotImplementedError
+ return imgs
+
+
+def add_gauss_blur(imgs, kernel_size=5, sigma=0):
+ outs = []
+ for k in range(imgs.shape[0]):
+ temp = imgs[k]
+ temp = cv2.GaussianBlur(temp, (kernel_size, kernel_size), sigma)
+ outs.append(temp)
+ outs = np.asarray(outs, dtype=np.float32)
+ outs[outs < 0] = 0
+ outs[outs > 1] = 1
+ return outs
+
+
+def add_sobel(imgs, if_mean=False):
+ outs = []
+ for k in range(imgs.shape[0]):
+ temp = imgs[k]
+ # sobelx = cv2.Sobel(temp, cv2.CV_32F, 1, 0)
+ # sobely = cv2.Sobel(temp, cv2.CV_32F, 0, 1)
+ # sobelx = filters.sobel_h(temp)
+ # sobely = filters.sobel_v(temp)
+ # dst = cv2.addWeighted(sobelx, 0.5, sobely, 0.5, 0)
+ # dst = sobelx * 0.5 + sobely * 0.5
+ # dst = cv2.Sobel(temp, cv2.CV_32F, 1, 1)
+ if if_mean:
+ mean = np.mean(temp)
+ else:
+ mean = 0
+ dst = filters.sobel(temp) + mean
+ outs.append(dst)
+ outs = np.asarray(outs, dtype=np.float32)
+ outs[outs < 0] = 0
+ outs[outs > 1] = 1
+ return outs
+
+
+def add_intensity(imgs, contrast_factor=0.1, brightness_factor=0.1):
+ # imgs *= 1 + (np.random.rand() - 0.5) * contrast_factor
+ # imgs += (np.random.rand() - 0.5) * brightness_factor
+ # imgs = np.clip(imgs, 0, 1)
+ # imgs **= 2.0**(np.random.rand()*2 - 1)
+ imgs *= 1 + contrast_factor
+ imgs += brightness_factor
+ imgs = np.clip(imgs, 0, 1)
+ return imgs
+
+
+def interp_5d(data, det_size, mode='bilinear'):
+ assert len(data.shape) == 5, "the dimension of data must be 5!"
+ out = []
+ depth = data.shape[2]
+ for k in range(depth):
+ temp = data[:, :, k, :, :]
+ if mode == 'bilinear':
+ temp = F.interpolate(temp, size=(det_size, det_size), mode='bilinear', align_corners=True)
+ elif mode == 'nearest':
+ temp = F.interpolate(temp, size=(det_size, det_size), mode='nearest')
+ out.append(temp)
+ out = torch.stack(out, dim=2)
+ return out
+
+
+def convert_consistency_scale(gt, det_size):
+ B, C, D, H, W = gt.shape
+ gt = gt.detach().clone()
+ out_gt = []
+ masks = []
+ for k in range(B):
+ gt_temp = gt[k]
+ det_size_temp = det_size[k]
+ if det_size_temp[0] == gt_temp.shape[-1]:
+ mask = torch.ones_like(gt_temp)
+ out_gt.append(gt_temp)
+ masks.append(mask)
+ elif det_size_temp[0] > gt_temp.shape[-1]:
+ shift = int((det_size_temp[0] - gt_temp.shape[-1]) // 2)
+ gt_padding = torch.zeros((1, C, D, int(det_size_temp[0]), int(det_size_temp[0]))).float().cuda()
+ mask = torch.zeros_like(gt_padding)
+ gt_padding[0, :, :, shift:-shift, shift:-shift] = gt_temp
+ mask[0, :, :, shift:-shift, shift:-shift] = 1
+ # gt_padding = F.interpolate(gt_padding, size=(D, int(gt_temp.shape[-1]), int(gt_temp.shape[-1])), mode='trilinear', align_corners=True)
+ gt_padding = interp_5d(gt_padding, int(gt_temp.shape[-1]), mode='bilinear')
+ mask = F.interpolate(mask, size=(D, int(gt_temp.shape[-1]), int(gt_temp.shape[-1])), mode='nearest')
+ gt_padding = torch.squeeze(gt_padding, dim=0)
+ mask = torch.squeeze(mask, dim=0)
+ out_gt.append(gt_padding)
+ masks.append(mask)
+ else:
+ shift = int((gt_temp.shape[-1] - det_size_temp[0]) // 2)
+ mask = torch.zeros_like(gt_temp)
+ mask[:, :, shift:-shift, shift:-shift] = 1
+ gt_padding = gt_temp[:, :, shift:-shift, shift:-shift]
+ gt_padding = gt_padding[None, ...]
+ # gt_padding = F.interpolate(gt_padding, size=(D, int(gt_temp.shape[-1]), int(gt_temp.shape[-1])), mode='trilinear', align_corners=True)
+ gt_padding = interp_5d(gt_padding, int(gt_temp.shape[-1]), mode='bilinear')
+ gt_padding = torch.squeeze(gt_padding, dim=0)
+ out_gt.append(gt_padding)
+ masks.append(mask)
+ out_gt = torch.stack(out_gt, dim=0)
+ masks = torch.stack(masks, dim=0)
+ return out_gt, masks
+
+
+def convert_consistency_flip(gt, rules):
+ B, C, D, H, W = gt.shape
+ gt = gt.detach().clone()
+ rules = rules.data.cpu().numpy()
+ out_gt = []
+ for k in range(B):
+ gt_temp = gt[k]
+ rule = rules[k]
+ gt_temp = simple_augment_torch(gt_temp, rule)
+ out_gt.append(gt_temp)
+ out_gt = torch.stack(out_gt, dim=0)
+ return out_gt
+
+
+class Rescale(object):
+ def __init__(self, scale_factor=2, det_shape=[18, 160, 160]):
+ super(Rescale, self).__init__()
+ self.scale_factor = scale_factor
+ self.det_shape = det_shape
+
+ def __call__(self, data):
+ src_shape = data.shape
+ assert src_shape[-1] >= self.det_shape[-1] * self.scale_factor, 'data shape must be 160*2'
+ min_size = self.det_shape[-1] // self.scale_factor
+ max_size = self.det_shape[-1] * self.scale_factor
+ scale_size = random.randint(min_size // 2, max_size // 2)
+ scale_size = scale_size * 2
+
+ if scale_size < src_shape[-1]:
+ shift = (src_shape[-1] - scale_size) // 2
+ data = data[:, shift:-shift, shift:-shift]
+ data = resize_3d(data, self.det_shape[-1], mode='linear')
+ return data, scale_size
+
+
+class Filp(object):
+ def __init__(self):
+ super(Filp, self).__init__()
+
+ def __call__(self, data):
+ rule = np.random.randint(2, size=4)
+ data = simple_augment(data, rule)
+ return data, rule
+
+
+# class Intensity(object):
+# def __init__(self, contrast_factor=0.3, brightness_factor=0.3):
+# super(Intensity, self).__init__()
+# self.CONTRAST_FACTOR = contrast_factor
+# self.BRIGHTNESS_FACTOR = brightness_factor
+
+# def __call__(self, data):
+# data = self._augment3D(data)
+# return data
+
+# def _augment3D(self, data, random_state=np.random):
+# """
+# Adapted from ELEKTRONN (http://elektronn.org/).
+# """
+# ran = random_state.rand(3)
+
+# transformedimgs = np.copy(data)
+# transformedimgs *= 1 + (ran[0] - 0.5)*self.CONTRAST_FACTOR
+# transformedimgs += (ran[1] - 0.5)*self.BRIGHTNESS_FACTOR
+# transformedimgs = np.clip(transformedimgs, 0, 1)
+# transformedimgs **= 2.0**(ran[2]*2 - 1)
+
+# return transformedimgs
+class Intensity(object):
+ def __init__(self, mode='mix',
+ skip_ratio=0.5,
+ CONTRAST_FACTOR=0.1,
+ BRIGHTNESS_FACTOR=0.1):
+ '''Image intensity augmentation, including adjusting contrast and brightness
+ Args:
+ mode: '2D', '3D' or 'mix' (contains '2D' and '3D')
+ skip_ratio: Probability of execution
+ CONTRAST_FACTOR: Contrast factor
+ BRIGHTNESS_FACTOR : Brightness factor
+ '''
+ super(Intensity, self).__init__()
+ assert mode == '3D' or mode == '2D' or mode == 'mix'
+ self.mode = mode
+ self.ratio = skip_ratio
+ self.CONTRAST_FACTOR = CONTRAST_FACTOR
+ self.BRIGHTNESS_FACTOR = BRIGHTNESS_FACTOR
+
+ def __call__(self, inputs):
+ return self.forward(inputs)
+
+ def forward(self, inputs):
+ inputs = inputs.copy()
+ skiprand = np.random.rand()
+ if self.mode == 'mix':
+ # The probability of '2D' is more than '3D'
+ threshold = 1 - (1 - self.ratio) / 2
+ mode_ = '3D' if skiprand > threshold else '2D'
+ else:
+ mode_ = self.mode
+ if mode_ == '2D':
+ inputs = self.augment2D(inputs)
+ elif mode_ == '3D':
+ inputs = self.augment3D(inputs)
+ inputs[inputs < 0] = 0
+ inputs[inputs > 1] = 1
+ return inputs
+
+ def augment2D(self, imgs):
+ for z in range(imgs.shape[-3]):
+ img = imgs[z, :, :]
+ img *= 1 + (np.random.rand() - 0.5) * self.CONTRAST_FACTOR
+ img += (np.random.rand() - 0.5) * self.BRIGHTNESS_FACTOR
+ img = np.clip(img, 0, 1)
+ img **= 2.0 ** (np.random.rand() * 2 - 1)
+ imgs[z, :, :] = img
+ return imgs
+
+ def augment3D(self, imgs):
+ imgs *= 1 + (np.random.rand() - 0.5) * self.CONTRAST_FACTOR
+ imgs += (np.random.rand() - 0.5) * self.BRIGHTNESS_FACTOR
+ imgs = np.clip(imgs, 0, 1)
+ imgs **= 2.0 ** (np.random.rand() * 2 - 1)
+ return imgs
+
+
+class GaussBlur(object):
+ def __init__(self, min_kernel=3, max_kernel=9, min_sigma=0, max_sigma=2):
+ super(GaussBlur, self).__init__()
+ self.min_kernel = min_kernel
+ self.max_kernel = max_kernel
+ self.min_sigma = min_sigma
+ self.max_sigma = max_sigma
+
+ def __call__(self, data):
+ kernel_size = random.randint(self.min_kernel // 2, self.max_kernel // 2)
+ kernel_size = kernel_size * 2 + 1
+ sigma = random.uniform(self.min_sigma, self.max_sigma)
+ data = add_gauss_blur(data, kernel_size=kernel_size, sigma=sigma)
+ return data
+
+
+class GaussNoise(object):
+ def __init__(self, min_std=0.01, max_std=0.2, norm_mode='trunc'):
+ super(GaussNoise, self).__init__()
+ self.min_std = min_std
+ self.max_std = max_std
+ self.norm_mode = norm_mode
+
+ def __call__(self, data):
+ std = random.uniform(self.min_std, self.max_std)
+ data = add_gauss_noise(data, std=std, norm_mode=self.norm_mode)
+ return data
+
+
+class Cutout(object):
+ def __init__(self, model_type='superhuman'):
+ super(Cutout, self).__init__()
+ self.model_type = model_type
+ # mask size
+ self.min_mask_size = [3, 5, 5]
+ self.max_mask_size = [5, 10, 10]
+ self.min_mask_counts = 20
+ self.max_mask_counts = 50
+ self.net_crop_size = [0, 0, 0]
+
+ def __call__(self, data):
+ mask_counts = random.randint(self.min_mask_counts, self.max_mask_counts)
+ mask_size_z = random.randint(self.min_mask_size[0], self.max_mask_size[0])
+ mask_size_xy = random.randint(self.min_mask_size[1], self.max_mask_size[1])
+ mask = gen_mask(data, net_crop_size=self.net_crop_size, \
+ mask_counts=mask_counts, \
+ mask_size_z=mask_size_z, \
+ mask_size_xy=mask_size_xy)
+ data = data * mask
+ return data
+
+
+class SobelFilter(object):
+ def __init__(self, if_mean=False):
+ super(SobelFilter, self).__init__()
+ self.if_mean = if_mean
+
+ def __call__(self, data):
+ data = add_sobel(data, if_mean=self.if_mean)
+ return data
+
+
+class Mixup(object):
+ def __init__(self, min_alpha=0.01, max_alpha=0.1):
+ super(Mixup, self).__init__()
+ self.min_alpha = min_alpha
+ self.max_alpha = max_alpha
+
+ def __call__(self, data, auxi):
+ alpha = random.uniform(self.min_alpha, self.max_alpha)
+ data = auxi * alpha + data * (1 - alpha)
+ data[data < 0] = 0
+ data[data > 1] = 1
+ return data
+
+
+class Missing(object):
+ '''Missing section augmentation
+ Args:
+ filling: the way of filling, 'zero' or 'random'
+ mode: 'mix', 'fully' or 'partially'
+ skip_ratio: Probability of execution
+ miss_ratio: Probability of missing
+ '''
+
+ def __init__(self, filling='zero', mode='mix', skip_ratio=0.5, miss_fully_ratio=0.2, miss_part_ratio=0.5):
+ super(Missing, self).__init__()
+ self.filling = filling
+ self.mode = mode
+ self.ratio = skip_ratio
+ self.miss_fully_ratio = miss_fully_ratio
+ self.miss_part_ratio = miss_part_ratio
+
+ def __call__(self, imgs):
+ return self.forward(imgs)
+
+ def forward(self, imgs):
+ imgs = imgs.copy()
+ if self.mode == 'mix':
+ r = np.random.rand()
+ mode_ = 'fully' if r < 0.5 else 'partially'
+ else:
+ mode_ = self.mode
+ if mode_ == 'fully':
+ imgs = self.augment_fully(imgs)
+ elif mode_ == 'partially':
+ imgs = self.augment_partially(imgs)
+ return imgs
+
+ def augment_fully(self, imgs):
+ d, h, w = imgs.shape
+ for i in range(d):
+ if np.random.rand() < self.miss_fully_ratio:
+ if self.filling == 'zero':
+ imgs[i] = 0
+ elif self.filling == 'random':
+ imgs[i] = np.random.rand(h, w)
+ return imgs
+
+ def augment_partially(self, imgs, size_ratio=0.3):
+ d, h, w = imgs.shape
+ for i in range(d):
+ if np.random.rand() < self.miss_part_ratio:
+ # randomly generate an area
+ sub_h = random.randint(int(h * size_ratio), int(h * (1 - size_ratio)))
+ sub_w = random.randint(int(w * size_ratio), int(w * (1 - size_ratio)))
+ start_h = random.randint(0, h - sub_h - 1)
+ start_w = random.randint(0, w - sub_w - 1)
+ if self.filling == 'zero':
+ imgs[i, start_h:start_h + sub_h, start_w:start_w + sub_w] = 0
+ elif self.filling == 'random':
+ imgs[i, start_h:start_h + sub_h, start_w:start_w + sub_w] = np.random.rand(sub_h, sub_w)
+ return imgs
+
+
+class BlurEnhanced(object):
+ '''Out-of-focus (Blur) section augmentation
+ Args:
+ mode: 'mix', 'fully' or 'partially'
+ skip_ratio: Probability of execution
+ blur_ratio: Probability of blur
+ '''
+
+ def __init__(self, mode='mix', skip_ratio=0.5, blur_fully_ratio=0.5, blur_part_ratio=0.7):
+ super(BlurEnhanced, self).__init__()
+ self.mode = mode
+ self.ratio = skip_ratio
+ self.blur_fully_ratio = blur_fully_ratio
+ self.blur_part_ratio = blur_part_ratio
+
+ def __call__(self, imgs):
+ return self.forward(imgs)
+
+ def forward(self, imgs):
+ imgs = imgs.copy()
+ if self.mode == 'mix':
+ r = np.random.rand()
+ mode_ = 'fully' if r < 0.5 else 'partially'
+ else:
+ mode_ = self.mode
+ if mode_ == 'fully':
+ imgs = self.augment_fully(imgs)
+ elif mode_ == 'partially':
+ imgs = self.augment_partially(imgs)
+ return imgs
+
+ def augment_fully(self, imgs):
+ d, h, w = imgs.shape
+ for i in range(d):
+ if np.random.rand() < self.blur_fully_ratio:
+ sigma = np.random.uniform(0, 5)
+ imgs[i] = gaussian_filter(imgs[i], sigma)
+ return imgs
+
+ def augment_partially(self, imgs, size_ratio=0.3):
+ d, h, w = imgs.shape
+ for i in range(d):
+ if np.random.rand() < self.blur_part_ratio:
+ # randomly generate an area
+ sub_h = random.randint(int(h * size_ratio), int(h * (1 - size_ratio)))
+ sub_w = random.randint(int(w * size_ratio), int(w * (1 - size_ratio)))
+ start_h = random.randint(0, h - sub_h - 1)
+ start_w = random.randint(0, w - sub_w - 1)
+ sigma = np.random.uniform(0, 5)
+ imgs[i, start_h:start_h + sub_h, start_w:start_w + sub_w] = \
+ gaussian_filter(imgs[i, start_h:start_h + sub_h, start_w:start_w + sub_w], sigma)
+ return imgs
+
+
+class Elastic(object):
+ '''Elasticly deform a batch. Requests larger batches upstream to avoid data
+ loss due to rotation and jitter.
+ Args:
+ control_point_spacing (``tuple`` of ``int``):
+ Distance between control points for the elastic deformation, in
+ voxels per dimension.
+ jitter_sigma (``tuple`` of ``float``):
+ Standard deviation of control point jitter distribution, in voxels
+ per dimension.
+ rotation_interval (``tuple`` of two ``floats``):
+ Interval to randomly sample rotation angles from (0, 2PI).
+ prob_slip (``float``):
+ Probability of a section to "slip", i.e., be independently moved in
+ x-y.
+ prob_shift (``float``):
+ Probability of a section and all following sections to move in x-y.
+ max_misalign (``int``):
+ Maximal voxels to shift in x and y. Samples will be drawn
+ uniformly. Used if ``prob_slip + prob_shift`` > 0.
+ subsample (``int``):
+ Instead of creating an elastic transformation on the full
+ resolution, create one subsampled by the given factor, and linearly
+ interpolate to obtain the full resolution transformation. This can
+ significantly speed up this node, at the expense of having visible
+ piecewise linear deformations for large factors. Usually, a factor
+ of 4 can savely by used without noticable changes. However, the
+ default is 1 (i.e., no subsampling).
+ '''
+
+ def __init__(
+ self,
+ control_point_spacing=[4, 40, 40],
+ jitter_sigma=[0, 0, 0], # recommend: [0, 2, 2]
+ rotation_interval=[0, 0],
+ prob_slip=0, # recommend: 0.05
+ prob_shift=0, # recommend: 0.05
+ max_misalign=0, # 17 in superhuman
+ subsample=1,
+ padding=None,
+ skip_ratio=0.5): # recommend: 10
+ super(Elastic, self).__init__()
+
+ self.control_point_spacing = control_point_spacing
+ self.jitter_sigma = jitter_sigma
+ self.rotation_start = rotation_interval[0]
+ self.rotation_max_amount = rotation_interval[1] - rotation_interval[0]
+ self.prob_slip = prob_slip
+ self.prob_shift = prob_shift
+ self.max_misalign = max_misalign
+ self.subsample = subsample
+ self.padding = padding
+ self.ratio = skip_ratio
+
+ def create_transformation(self, target_shape):
+ transformation = create_identity_transformation(
+ target_shape,
+ subsample=self.subsample)
+ # shape: channel,d,w,h
+
+ # elastic ##cost time##
+ if sum(self.jitter_sigma) > 0:
+ transformation += create_elastic_transformation(
+ target_shape,
+ self.control_point_spacing,
+ self.jitter_sigma,
+ subsample=self.subsample)
+
+ # rotation = random.random()*self.rotation_max_amount + self.rotation_start
+ # if rotation != 0:
+ # transformation += create_rotation_transformation(
+ # target_shape,
+ # rotation,
+ # subsample=self.subsample)
+
+ # if self.subsample > 1:
+ # transformation = upscale_transformation(
+ # transformation,
+ # tuple(target_shape))
+
+ if self.prob_slip + self.prob_shift > 0:
+ misalign(transformation, self.prob_slip,
+ self.prob_shift, self.max_misalign)
+
+ return transformation
+
+ def __call__(self, imgs):
+ return self.forward(imgs)
+
+ def forward(self, imgs):
+ '''Args:
+ imgs: numpy array, [Z, Y, Z], it always is float and 0~1
+ mask: numpy array, [Z, Y, Z], it always is uint16
+ '''
+ imgs = imgs.copy()
+ if self.padding is not None:
+ imgs = np.pad(imgs, ((0, 0), \
+ (self.padding, self.padding), \
+ (self.padding, self.padding)), mode='reflect')
+ transform = self.create_transformation(imgs.shape)
+ img_transform = apply_transformation(imgs,
+ transform,
+ interpolate=False,
+ outside_value=0, # imgs.dtype.type(-1)
+ output=np.zeros(imgs.shape, dtype=np.float32))
+ # seg_transform[seg_transform < 0] = 0
+ # seg_transform[seg_transform > 60000] = 0
+ if self.padding is not None and self.padding != 0:
+ img_transform = img_transform[:, self.padding:-self.padding, self.padding:-self.padding]
+ return img_transform
+
+
+class Artifact(object):
+ def __init__(self, min_sec=1, max_sec=5):
+ super(Artifact, self).__init__()
+ self.min_sec = min_sec
+ self.max_sec = max_sec
+ self.offset = 40
+
+ def __call__(self, data):
+ data = data.copy()
+ num_sec = random.randint(self.min_sec, self.max_sec)
+ num_imgs = data.shape[0]
+ rand_sample = random.sample(range(num_imgs), num_sec)
+ for k in rand_sample:
+ tmp = data[k].copy()
+ tmp = (tmp * 255).astype(np.uint8)
+ tmp = self.degradation(tmp)
+ data[k] = tmp.astype(np.float32) / 255.0
+ return data
+
+ def degradation(self, img):
+ img = np.pad(img, ((self.offset, self.offset), (self.offset, self.offset)), mode='reflect')
+ height, width = img.shape
+ line_width = random.randint(5, 10)
+ fold_width = random.randint(line_width + 1, 40)
+
+ # two end points
+ # 1 --> top line (0, x)
+ # 2 --> right line (x, width)
+ # 3 --> bottom line (height, x)
+ # 4 --> left line (x, 0)
+ k1 = random.randint(1, 4)
+ k2 = random.randint(1, 4)
+ while k1 == k2:
+ k2 = random.randint(1, 4)
+
+ if k1 == 1:
+ x = random.randint(1, width - 1)
+ p1 = [0, x]
+ elif k1 == 2:
+ x = random.randint(1, height - 1)
+ p1 = [x, width]
+ elif k1 == 3:
+ x = random.randint(1, width - 1)
+ p1 = [height, x]
+ else:
+ x = random.randint(1, height - 1)
+ p1 = [x, 0]
+
+ if k2 == 1:
+ x = random.randint(1, width - 1)
+ p2 = [0, x]
+ elif k2 == 2:
+ x = random.randint(1, height - 1)
+ p2 = [x, width]
+ elif k2 == 3:
+ x = random.randint(1, width - 1)
+ p2 = [height, x]
+ else:
+ x = random.randint(1, height - 1)
+ p2 = [x, 0]
+
+ dis_k = random.uniform(0.00001, 0.1)
+ k, b = gen_line(p1, p2)
+ flow, flow2, mask = gen_flow(height, width, k, b, line_width, fold_width, dis_k)
+
+ deformed = image_warp(img, flow, mode='bilinear') # nearest or bilinear
+ deformed = (deformed * mask).astype(np.uint8)
+ deformed = deformed[self.offset:-self.offset, self.offset:-self.offset]
+
+ return deformed
diff --git a/legacy/Train_and_Inference/utils/consistency_aug_perturbations_sup.py b/legacy/Train_and_Inference/utils/consistency_aug_perturbations_sup.py
new file mode 100644
index 0000000..346622e
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/consistency_aug_perturbations_sup.py
@@ -0,0 +1,763 @@
+import cv2
+import torch
+import random
+import numpy as np
+import torch.nn.functional as F
+from skimage import filters
+from scipy.ndimage.filters import gaussian_filter
+
+from utils.augmentation import create_identity_transformation
+from utils.augmentation import create_elastic_transformation
+from utils.augmentation import apply_transformation
+from utils.augmentation import misalign
+from utils.flow_synthesis import gen_line, gen_flow
+from utils.image_warp import image_warp
+
+
+def simple_augment(data, rule):
+ assert np.size(rule) == 4
+ assert data.ndim == 3
+ # z reflection
+ if rule[0]:
+ data = data[::-1, :, :]
+ # x reflection
+ if rule[1]:
+ data = data[:, :, ::-1]
+ # y reflection
+ if rule[2]:
+ data = data[:, ::-1, :]
+ # transpose in xy
+ if rule[3]:
+ data = np.transpose(data, (0, 2, 1))
+ return data
+
+
+def simple_augment_torch(data, rule):
+ assert np.size(rule) == 4
+ assert len(data.shape) == 4
+ # z reflection
+ if rule[0]:
+ data = torch.flip(data, [1])
+ # x reflection
+ if rule[1]:
+ data = torch.flip(data, [3])
+ # y reflection
+ if rule[2]:
+ data = torch.flip(data, [2])
+ # transpose in xy
+ if rule[3]:
+ data = data.permute(0, 1, 3, 2)
+ return data
+
+
+def simple_augment_reverse(data, rule):
+ assert np.size(rule) == 4
+ assert len(data.shape) == 5
+ # transpose in xy
+ if rule[3]:
+ # data = np.transpose(data, (0, 1, 2, 4, 3))
+ data = data.permute(0, 1, 2, 4, 3)
+ # y reflection
+ if rule[2]:
+ # data = data[:, :, :, ::-1, :]
+ data = torch.flip(data, [3])
+ # x reflection
+ if rule[1]:
+ # data = data[:, :, :, :, ::-1]
+ data = torch.flip(data, [4])
+ # z reflection
+ if rule[0]:
+ # data = data[:, :, ::-1, :, :]
+ data = torch.flip(data, [2])
+ return data
+
+
+def order_aug(imgs, num_patch=4):
+ assert imgs.shape[-1] % num_patch == 0
+ patch_size = imgs.shape[-1] // num_patch
+ new_imgs = np.zeros_like(imgs, dtype=np.float32)
+ # ran_order = np.random.shuffle(np.arange(num_patch**2))
+ ran_order = np.random.permutation(num_patch ** 2)
+ for k in range(num_patch ** 2):
+ xid_new = k // num_patch
+ yid_new = k % num_patch
+ order_id = ran_order[k]
+ xid_old = order_id // num_patch
+ yid_old = order_id % num_patch
+ new_imgs[:, xid_new * patch_size:(xid_new + 1) * patch_size, yid_new * patch_size:(yid_new + 1) * patch_size] = \
+ imgs[:, xid_old * patch_size:(xid_old + 1) * patch_size, yid_old * patch_size:(yid_old + 1) * patch_size]
+ return new_imgs
+
+
+def gen_mask(imgs, net_crop_size=[0, 0, 0], mask_counts=80, mask_size_z=8, mask_size_xy=15):
+ crop_size = list(imgs.shape)
+ mask = np.ones_like(imgs, dtype=np.float32)
+ for k in range(mask_counts):
+ mz = random.randint(net_crop_size[0], crop_size[0] - mask_size_z - net_crop_size[0])
+ my = random.randint(net_crop_size[1], crop_size[1] - mask_size_xy - net_crop_size[1])
+ mx = random.randint(net_crop_size[2], crop_size[2] - mask_size_xy - net_crop_size[2])
+ mask[mz:mz + mask_size_z, my:my + mask_size_xy, mx:mx + mask_size_xy] = 0
+ return mask
+
+
+def resize_3d(imgs, det_size, mode='linear'):
+ new_imgs = []
+ for k in range(imgs.shape[0]):
+ temp = imgs[k]
+ if mode == 'linear':
+ temp = cv2.resize(temp, (det_size, det_size), interpolation=cv2.INTER_LINEAR)
+ elif mode == 'nearest':
+ temp = cv2.resize(temp, (det_size, det_size), interpolation=cv2.INTER_NEAREST)
+ else:
+ raise AttributeError('No this interpolation mode!')
+ new_imgs.append(temp)
+ new_imgs = np.asarray(new_imgs)
+ return new_imgs
+
+
+def add_gauss_noise(imgs, std=0.01, norm_mode='norm'):
+ gaussian = np.random.normal(0, std, (imgs.shape))
+ imgs = imgs + gaussian
+ if norm_mode == 'norm':
+ imgs = (imgs - np.min(imgs)) / (np.max(imgs) - np.min(imgs))
+ elif norm_mode == 'trunc':
+ imgs[imgs < 0] = 0
+ imgs[imgs > 1] = 1
+ else:
+ raise NotImplementedError
+ return imgs
+
+
+def add_gauss_blur(imgs, kernel_size=5, sigma=0):
+ outs = []
+ for k in range(imgs.shape[0]):
+ temp = imgs[k]
+ temp = cv2.GaussianBlur(temp, (kernel_size, kernel_size), sigma)
+ outs.append(temp)
+ outs = np.asarray(outs, dtype=np.float32)
+ outs[outs < 0] = 0
+ outs[outs > 1] = 1
+ return outs
+
+
+def add_sobel(imgs, if_mean=False):
+ outs = []
+ for k in range(imgs.shape[0]):
+ temp = imgs[k]
+ # sobelx = cv2.Sobel(temp, cv2.CV_32F, 1, 0)
+ # sobely = cv2.Sobel(temp, cv2.CV_32F, 0, 1)
+ # sobelx = filters.sobel_h(temp)
+ # sobely = filters.sobel_v(temp)
+ # dst = cv2.addWeighted(sobelx, 0.5, sobely, 0.5, 0)
+ # dst = sobelx * 0.5 + sobely * 0.5
+ # dst = cv2.Sobel(temp, cv2.CV_32F, 1, 1)
+ if if_mean:
+ mean = np.mean(temp)
+ else:
+ mean = 0
+ dst = filters.sobel(temp) + mean
+ outs.append(dst)
+ outs = np.asarray(outs, dtype=np.float32)
+ outs[outs < 0] = 0
+ outs[outs > 1] = 1
+ return outs
+
+
+def add_intensity(imgs, contrast_factor=0.1, brightness_factor=0.1):
+ # imgs *= 1 + (np.random.rand() - 0.5) * contrast_factor
+ # imgs += (np.random.rand() - 0.5) * brightness_factor
+ # imgs = np.clip(imgs, 0, 1)
+ # imgs **= 2.0**(np.random.rand()*2 - 1)
+ imgs *= 1 + contrast_factor
+ imgs += brightness_factor
+ imgs = np.clip(imgs, 0, 1)
+ return imgs
+
+
+def interp_5d(data, det_size, mode='bilinear'):
+ assert len(data.shape) == 5, "the dimension of data must be 5!"
+ out = []
+ depth = data.shape[2]
+ for k in range(depth):
+ temp = data[:, :, k, :, :]
+ if mode == 'bilinear':
+ temp = F.interpolate(temp, size=(det_size, det_size), mode='bilinear', align_corners=True)
+ elif mode == 'nearest':
+ temp = F.interpolate(temp, size=(det_size, det_size), mode='nearest')
+ out.append(temp)
+ out = torch.stack(out, dim=2)
+ return out
+
+
+def convert_consistency_scale(gt, det_size):
+ B, C, D, H, W = gt.shape
+ gt = gt.detach().clone()
+ out_gt = []
+ masks = []
+ for k in range(B):
+ gt_temp = gt[k]
+ det_size_temp = det_size[k]
+ if det_size_temp[0] == gt_temp.shape[-1]:
+ mask = torch.ones_like(gt_temp)
+ out_gt.append(gt_temp)
+ masks.append(mask)
+ elif det_size_temp[0] > gt_temp.shape[-1]:
+ shift = int((det_size_temp[0] - gt_temp.shape[-1]) // 2)
+ gt_padding = torch.zeros((1, C, D, int(det_size_temp[0]), int(det_size_temp[0]))).float().cuda()
+ mask = torch.zeros_like(gt_padding)
+ gt_padding[0, :, :, shift:-shift, shift:-shift] = gt_temp
+ mask[0, :, :, shift:-shift, shift:-shift] = 1
+ # gt_padding = F.interpolate(gt_padding, size=(D, int(gt_temp.shape[-1]), int(gt_temp.shape[-1])), mode='trilinear', align_corners=True)
+ gt_padding = interp_5d(gt_padding, int(gt_temp.shape[-1]), mode='bilinear')
+ mask = F.interpolate(mask, size=(D, int(gt_temp.shape[-1]), int(gt_temp.shape[-1])), mode='nearest')
+ gt_padding = torch.squeeze(gt_padding, dim=0)
+ mask = torch.squeeze(mask, dim=0)
+ out_gt.append(gt_padding)
+ masks.append(mask)
+ else:
+ shift = int((gt_temp.shape[-1] - det_size_temp[0]) // 2)
+ mask = torch.zeros_like(gt_temp)
+ mask[:, :, shift:-shift, shift:-shift] = 1
+ gt_padding = gt_temp[:, :, shift:-shift, shift:-shift]
+ gt_padding = gt_padding[None, ...]
+ # gt_padding = F.interpolate(gt_padding, size=(D, int(gt_temp.shape[-1]), int(gt_temp.shape[-1])), mode='trilinear', align_corners=True)
+ gt_padding = interp_5d(gt_padding, int(gt_temp.shape[-1]), mode='bilinear')
+ gt_padding = torch.squeeze(gt_padding, dim=0)
+ out_gt.append(gt_padding)
+ masks.append(mask)
+ out_gt = torch.stack(out_gt, dim=0)
+ masks = torch.stack(masks, dim=0)
+ return out_gt, masks
+
+
+def convert_consistency_flip(gt, rules):
+ B, C, D, H, W = gt.shape
+ gt = gt.detach().clone()
+ rules = rules.data.cpu().numpy()
+ out_gt = []
+ for k in range(B):
+ gt_temp = gt[k]
+ rule = rules[k]
+ gt_temp = simple_augment_torch(gt_temp, rule)
+ out_gt.append(gt_temp)
+ out_gt = torch.stack(out_gt, dim=0)
+ return out_gt
+
+
+class Rescale(object):
+ def __init__(self, scale_factor=2, det_shape=[18, 160, 160]):
+ super(Rescale, self).__init__()
+ self.scale_factor = scale_factor
+ self.det_shape = det_shape
+
+ def __call__(self, data):
+ src_shape = data.shape
+ assert src_shape[-1] >= self.det_shape[-1] * self.scale_factor, 'data shape must be 160*2'
+ min_size = self.det_shape[-1] // self.scale_factor
+ max_size = self.det_shape[-1] * self.scale_factor
+ scale_size = random.randint(min_size // 2, max_size // 2)
+ scale_size = scale_size * 2
+
+ if scale_size < src_shape[-1]:
+ shift = (src_shape[-1] - scale_size) // 2
+ data = data[:, shift:-shift, shift:-shift]
+ data = resize_3d(data, self.det_shape[-1], mode='linear')
+ return data, scale_size
+
+
+class Filp(object):
+ def __init__(self):
+ super(Filp, self).__init__()
+
+ def __call__(self, data):
+ rule = np.random.randint(2, size=4)
+ data = simple_augment(data, rule)
+ return data, rule
+
+
+# class Intensity(object):
+# def __init__(self, contrast_factor=0.3, brightness_factor=0.3):
+# super(Intensity, self).__init__()
+# self.CONTRAST_FACTOR = contrast_factor
+# self.BRIGHTNESS_FACTOR = brightness_factor
+
+# def __call__(self, data):
+# data = self._augment3D(data)
+# return data
+
+# def _augment3D(self, data, random_state=np.random):
+# """
+# Adapted from ELEKTRONN (http://elektronn.org/).
+# """
+# ran = random_state.rand(3)
+
+# transformedimgs = np.copy(data)
+# transformedimgs *= 1 + (ran[0] - 0.5)*self.CONTRAST_FACTOR
+# transformedimgs += (ran[1] - 0.5)*self.BRIGHTNESS_FACTOR
+# transformedimgs = np.clip(transformedimgs, 0, 1)
+# transformedimgs **= 2.0**(ran[2]*2 - 1)
+
+# return transformedimgs
+class Intensity(object):
+ def __init__(self, mode='mix',
+ skip_ratio=0.5,
+ CONTRAST_FACTOR=0.1,
+ BRIGHTNESS_FACTOR=0.1):
+ '''Image intensity augmentation, including adjusting contrast and brightness
+ Args:
+ mode: '2D', '3D' or 'mix' (contains '2D' and '3D')
+ skip_ratio: Probability of execution
+ CONTRAST_FACTOR: Contrast factor
+ BRIGHTNESS_FACTOR : Brightness factor
+ '''
+ super(Intensity, self).__init__()
+ assert mode == '3D' or mode == '2D' or mode == 'mix'
+ self.mode = mode
+ self.ratio = skip_ratio
+ self.CONTRAST_FACTOR = CONTRAST_FACTOR
+ self.BRIGHTNESS_FACTOR = BRIGHTNESS_FACTOR
+
+ def __call__(self, inputs):
+ return self.forward(inputs)
+
+ def forward(self, inputs):
+ inputs = inputs.copy()
+ skiprand = np.random.rand()
+ if self.mode == 'mix':
+ # The probability of '2D' is more than '3D'
+ threshold = 1 - (1 - self.ratio) / 2
+ mode_ = '3D' if skiprand > threshold else '2D'
+ else:
+ mode_ = self.mode
+ if mode_ == '2D':
+ inputs = self.augment2D(inputs)
+ elif mode_ == '3D':
+ inputs = self.augment3D(inputs)
+ inputs[inputs < 0] = 0
+ inputs[inputs > 1] = 1
+ return inputs
+
+ def augment2D(self, imgs):
+ for z in range(imgs.shape[-3]):
+ img = imgs[z, :, :]
+ img *= 1 + (np.random.rand() - 0.5) * self.CONTRAST_FACTOR
+ img += (np.random.rand() - 0.5) * self.BRIGHTNESS_FACTOR
+ img = np.clip(img, 0, 1)
+ img **= 2.0 ** (np.random.rand() * 2 - 1)
+ imgs[z, :, :] = img
+ return imgs
+
+ def augment3D(self, imgs):
+ imgs *= 1 + (np.random.rand() - 0.5) * self.CONTRAST_FACTOR
+ imgs += (np.random.rand() - 0.5) * self.BRIGHTNESS_FACTOR
+ imgs = np.clip(imgs, 0, 1)
+ imgs **= 2.0 ** (np.random.rand() * 2 - 1)
+ return imgs
+
+
+class GaussBlur(object):
+ def __init__(self, min_kernel=3, max_kernel=9, min_sigma=0, max_sigma=2):
+ super(GaussBlur, self).__init__()
+ self.min_kernel = min_kernel
+ self.max_kernel = max_kernel
+ self.min_sigma = min_sigma
+ self.max_sigma = max_sigma
+
+ def __call__(self, data):
+ kernel_size = random.randint(self.min_kernel // 2, self.max_kernel // 2)
+ kernel_size = kernel_size * 2 + 1
+ sigma = random.uniform(self.min_sigma, self.max_sigma)
+ data = add_gauss_blur(data, kernel_size=kernel_size, sigma=sigma)
+ return data
+
+
+class GaussNoise(object):
+ def __init__(self, min_std=0.01, max_std=0.2, norm_mode='trunc'):
+ super(GaussNoise, self).__init__()
+ self.min_std = min_std
+ self.max_std = max_std
+ self.norm_mode = norm_mode
+
+ def __call__(self, data):
+ std = random.uniform(self.min_std, self.max_std)
+ data = add_gauss_noise(data, std=std, norm_mode=self.norm_mode)
+ return data
+
+
+class Cutout(object):
+ def __init__(self, model_type='superhuman'):
+ super(Cutout, self).__init__()
+ self.model_type = model_type
+ # mask size
+ if self.model_type == 'mala':
+ self.min_mask_size = [5, 5, 5]
+ self.max_mask_size = [8, 12, 12]
+ self.min_mask_counts = 40
+ self.max_mask_counts = 60
+ self.net_crop_size = [14, 106, 106]
+ else:
+ self.min_mask_size = [5, 15, 15]
+ self.max_mask_size = [10, 25, 25]
+ self.min_mask_counts = 20
+ self.max_mask_counts = 50
+ self.net_crop_size = [0, 0, 0]
+
+ def __call__(self, data):
+ mask_counts = random.randint(self.min_mask_counts, self.max_mask_counts)
+ mask_size_z = random.randint(self.min_mask_size[0], self.max_mask_size[0])
+ mask_size_xy = random.randint(self.min_mask_size[1], self.max_mask_size[1])
+ mask = gen_mask(data, net_crop_size=self.net_crop_size, \
+ mask_counts=mask_counts, \
+ mask_size_z=mask_size_z, \
+ mask_size_xy=mask_size_xy)
+ data = data * mask
+ return data
+
+class Cutout_P(object):
+ def __init__(self, model_type='superhuman'):
+ super(Cutout, self).__init__()
+ self.model_type = model_type
+ # mask size
+ if self.model_type == 'mala':
+ self.min_mask_size = [5, 5, 5]
+ self.max_mask_size = [8, 12, 12]
+ self.min_mask_counts = 40
+ self.max_mask_counts = 60
+ self.net_crop_size = [14, 106, 106]
+ else:
+ self.min_mask_size = [3, 9, 9]
+ self.max_mask_size = [5, 15, 15]
+ self.min_mask_counts = 10
+ self.max_mask_counts = 30
+ self.net_crop_size = [0, 0, 0]
+
+ def __call__(self, data):
+ mask_counts = random.randint(self.min_mask_counts, self.max_mask_counts)
+ mask_size_z = random.randint(self.min_mask_size[0], self.max_mask_size[0])
+ mask_size_xy = random.randint(self.min_mask_size[1], self.max_mask_size[1])
+ mask = gen_mask(data, net_crop_size=self.net_crop_size, \
+ mask_counts=mask_counts, \
+ mask_size_z=mask_size_z, \
+ mask_size_xy=mask_size_xy)
+ data = data * mask
+ return data
+
+class SobelFilter(object):
+ def __init__(self, if_mean=False):
+ super(SobelFilter, self).__init__()
+ self.if_mean = if_mean
+
+ def __call__(self, data):
+ data = add_sobel(data, if_mean=self.if_mean)
+ return data
+
+
+class Mixup(object):
+ def __init__(self, min_alpha=0.01, max_alpha=0.1):
+ super(Mixup, self).__init__()
+ self.min_alpha = min_alpha
+ self.max_alpha = max_alpha
+
+ def __call__(self, data, auxi):
+ alpha = random.uniform(self.min_alpha, self.max_alpha)
+ data = auxi * alpha + data * (1 - alpha)
+ data[data < 0] = 0
+ data[data > 1] = 1
+ return data
+
+
+class Missing(object):
+ '''Missing section augmentation
+ Args:
+ filling: the way of filling, 'zero' or 'random'
+ mode: 'mix', 'fully' or 'partially'
+ skip_ratio: Probability of execution
+ miss_ratio: Probability of missing
+ '''
+
+ def __init__(self, filling='zero', mode='mix', skip_ratio=0.5, miss_fully_ratio=0.2, miss_part_ratio=0.5):
+ super(Missing, self).__init__()
+ self.filling = filling
+ self.mode = mode
+ self.ratio = skip_ratio
+ self.miss_fully_ratio = miss_fully_ratio
+ self.miss_part_ratio = miss_part_ratio
+
+ def __call__(self, imgs):
+ return self.forward(imgs)
+
+ def forward(self, imgs):
+ imgs = imgs.copy()
+ if self.mode == 'mix':
+ r = np.random.rand()
+ mode_ = 'fully' if r < 0.5 else 'partially'
+ else:
+ mode_ = self.mode
+ if mode_ == 'fully':
+ imgs = self.augment_fully(imgs)
+ elif mode_ == 'partially':
+ imgs = self.augment_partially(imgs)
+ return imgs
+
+ def augment_fully(self, imgs):
+ d, h, w = imgs.shape
+ for i in range(d):
+ if np.random.rand() < self.miss_fully_ratio:
+ if self.filling == 'zero':
+ imgs[i] = 0
+ elif self.filling == 'random':
+ imgs[i] = np.random.rand(h, w)
+ return imgs
+
+ def augment_partially(self, imgs, size_ratio=0.3):
+ d, h, w = imgs.shape
+ for i in range(d):
+ if np.random.rand() < self.miss_part_ratio:
+ # randomly generate an area
+ sub_h = random.randint(int(h * size_ratio), int(h * (1 - size_ratio)))
+ sub_w = random.randint(int(w * size_ratio), int(w * (1 - size_ratio)))
+ start_h = random.randint(0, h - sub_h - 1)
+ start_w = random.randint(0, w - sub_w - 1)
+ if self.filling == 'zero':
+ imgs[i, start_h:start_h + sub_h, start_w:start_w + sub_w] = 0
+ elif self.filling == 'random':
+ imgs[i, start_h:start_h + sub_h, start_w:start_w + sub_w] = np.random.rand(sub_h, sub_w)
+ return imgs
+
+
+class BlurEnhanced(object):
+ '''Out-of-focus (Blur) section augmentation
+ Args:
+ mode: 'mix', 'fully' or 'partially'
+ skip_ratio: Probability of execution
+ blur_ratio: Probability of blur
+ '''
+
+ def __init__(self, mode='mix', skip_ratio=0.5, blur_fully_ratio=0.5, blur_part_ratio=0.7):
+ super(BlurEnhanced, self).__init__()
+ self.mode = mode
+ self.ratio = skip_ratio
+ self.blur_fully_ratio = blur_fully_ratio
+ self.blur_part_ratio = blur_part_ratio
+
+ def __call__(self, imgs):
+ return self.forward(imgs)
+
+ def forward(self, imgs):
+ imgs = imgs.copy()
+ if self.mode == 'mix':
+ r = np.random.rand()
+ mode_ = 'fully' if r < 0.5 else 'partially'
+ else:
+ mode_ = self.mode
+ if mode_ == 'fully':
+ imgs = self.augment_fully(imgs)
+ elif mode_ == 'partially':
+ imgs = self.augment_partially(imgs)
+ return imgs
+
+ def augment_fully(self, imgs):
+ d, h, w = imgs.shape
+ for i in range(d):
+ if np.random.rand() < self.blur_fully_ratio:
+ sigma = np.random.uniform(0, 5)
+ imgs[i] = gaussian_filter(imgs[i], sigma)
+ return imgs
+
+ def augment_partially(self, imgs, size_ratio=0.3):
+ d, h, w = imgs.shape
+ for i in range(d):
+ if np.random.rand() < self.blur_part_ratio:
+ # randomly generate an area
+ sub_h = random.randint(int(h * size_ratio), int(h * (1 - size_ratio)))
+ sub_w = random.randint(int(w * size_ratio), int(w * (1 - size_ratio)))
+ start_h = random.randint(0, h - sub_h - 1)
+ start_w = random.randint(0, w - sub_w - 1)
+ sigma = np.random.uniform(0, 5)
+ imgs[i, start_h:start_h + sub_h, start_w:start_w + sub_w] = \
+ gaussian_filter(imgs[i, start_h:start_h + sub_h, start_w:start_w + sub_w], sigma)
+ return imgs
+
+
+class Elastic(object):
+ '''Elasticly deform a batch. Requests larger batches upstream to avoid data
+ loss due to rotation and jitter.
+ Args:
+ control_point_spacing (``tuple`` of ``int``):
+ Distance between control points for the elastic deformation, in
+ voxels per dimension.
+ jitter_sigma (``tuple`` of ``float``):
+ Standard deviation of control point jitter distribution, in voxels
+ per dimension.
+ rotation_interval (``tuple`` of two ``floats``):
+ Interval to randomly sample rotation angles from (0, 2PI).
+ prob_slip (``float``):
+ Probability of a section to "slip", i.e., be independently moved in
+ x-y.
+ prob_shift (``float``):
+ Probability of a section and all following sections to move in x-y.
+ max_misalign (``int``):
+ Maximal voxels to shift in x and y. Samples will be drawn
+ uniformly. Used if ``prob_slip + prob_shift`` > 0.
+ subsample (``int``):
+ Instead of creating an elastic transformation on the full
+ resolution, create one subsampled by the given factor, and linearly
+ interpolate to obtain the full resolution transformation. This can
+ significantly speed up this node, at the expense of having visible
+ piecewise linear deformations for large factors. Usually, a factor
+ of 4 can savely by used without noticable changes. However, the
+ default is 1 (i.e., no subsampling).
+ '''
+
+ def __init__(
+ self,
+ control_point_spacing=[4, 40, 40],
+ jitter_sigma=[0, 0, 0], # recommend: [0, 2, 2]
+ rotation_interval=[0, 0],
+ prob_slip=0, # recommend: 0.05
+ prob_shift=0, # recommend: 0.05
+ max_misalign=0, # 17 in superhuman
+ subsample=1,
+ padding=None,
+ skip_ratio=0.5): # recommend: 10
+ super(Elastic, self).__init__()
+
+ self.control_point_spacing = control_point_spacing
+ self.jitter_sigma = jitter_sigma
+ self.rotation_start = rotation_interval[0]
+ self.rotation_max_amount = rotation_interval[1] - rotation_interval[0]
+ self.prob_slip = prob_slip
+ self.prob_shift = prob_shift
+ self.max_misalign = max_misalign
+ self.subsample = subsample
+ self.padding = padding
+ self.ratio = skip_ratio
+
+ def create_transformation(self, target_shape):
+ transformation = create_identity_transformation(
+ target_shape,
+ subsample=self.subsample)
+ # shape: channel,d,w,h
+
+ # elastic ##cost time##
+ if sum(self.jitter_sigma) > 0:
+ transformation += create_elastic_transformation(
+ target_shape,
+ self.control_point_spacing,
+ self.jitter_sigma,
+ subsample=self.subsample)
+
+ # rotation = random.random()*self.rotation_max_amount + self.rotation_start
+ # if rotation != 0:
+ # transformation += create_rotation_transformation(
+ # target_shape,
+ # rotation,
+ # subsample=self.subsample)
+
+ # if self.subsample > 1:
+ # transformation = upscale_transformation(
+ # transformation,
+ # tuple(target_shape))
+
+ if self.prob_slip + self.prob_shift > 0:
+ misalign(transformation, self.prob_slip,
+ self.prob_shift, self.max_misalign)
+
+ return transformation
+
+ def __call__(self, imgs):
+ return self.forward(imgs)
+
+ def forward(self, imgs):
+ '''Args:
+ imgs: numpy array, [Z, Y, Z], it always is float and 0~1
+ mask: numpy array, [Z, Y, Z], it always is uint16
+ '''
+ imgs = imgs.copy()
+ if self.padding is not None:
+ imgs = np.pad(imgs, ((0, 0), \
+ (self.padding, self.padding), \
+ (self.padding, self.padding)), mode='reflect')
+ transform = self.create_transformation(imgs.shape)
+ img_transform = apply_transformation(imgs,
+ transform,
+ interpolate=False,
+ outside_value=0, # imgs.dtype.type(-1)
+ output=np.zeros(imgs.shape, dtype=np.float32))
+ # seg_transform[seg_transform < 0] = 0
+ # seg_transform[seg_transform > 60000] = 0
+ if self.padding is not None and self.padding != 0:
+ img_transform = img_transform[:, self.padding:-self.padding, self.padding:-self.padding]
+ return img_transform
+
+
+class Artifact(object):
+ def __init__(self, min_sec=1, max_sec=5):
+ super(Artifact, self).__init__()
+ self.min_sec = min_sec
+ self.max_sec = max_sec
+ self.offset = 40
+
+ def __call__(self, data):
+ data = data.copy()
+ num_sec = random.randint(self.min_sec, self.max_sec)
+ num_imgs = data.shape[0]
+ rand_sample = random.sample(range(num_imgs), num_sec)
+ for k in rand_sample:
+ tmp = data[k].copy()
+ tmp = (tmp * 255).astype(np.uint8)
+ tmp = self.degradation(tmp)
+ data[k] = tmp.astype(np.float32) / 255.0
+ return data
+
+ def degradation(self, img):
+ img = np.pad(img, ((self.offset, self.offset), (self.offset, self.offset)), mode='reflect')
+ height, width = img.shape
+ line_width = random.randint(5, 10)
+ fold_width = random.randint(line_width + 1, 40)
+
+ # two end points
+ # 1 --> top line (0, x)
+ # 2 --> right line (x, width)
+ # 3 --> bottom line (height, x)
+ # 4 --> left line (x, 0)
+ k1 = random.randint(1, 4)
+ k2 = random.randint(1, 4)
+ while k1 == k2:
+ k2 = random.randint(1, 4)
+
+ if k1 == 1:
+ x = random.randint(1, width - 1)
+ p1 = [0, x]
+ elif k1 == 2:
+ x = random.randint(1, height - 1)
+ p1 = [x, width]
+ elif k1 == 3:
+ x = random.randint(1, width - 1)
+ p1 = [height, x]
+ else:
+ x = random.randint(1, height - 1)
+ p1 = [x, 0]
+
+ if k2 == 1:
+ x = random.randint(1, width - 1)
+ p2 = [0, x]
+ elif k2 == 2:
+ x = random.randint(1, height - 1)
+ p2 = [x, width]
+ elif k2 == 3:
+ x = random.randint(1, width - 1)
+ p2 = [height, x]
+ else:
+ x = random.randint(1, height - 1)
+ p2 = [x, 0]
+
+ dis_k = random.uniform(0.00001, 0.1)
+ k, b = gen_line(p1, p2)
+ flow, flow2, mask = gen_flow(height, width, k, b, line_width, fold_width, dis_k)
+
+ deformed = image_warp(img, flow, mode='bilinear') # nearest or bilinear
+ deformed = (deformed * mask).astype(np.uint8)
+ deformed = deformed[self.offset:-self.offset, self.offset:-self.offset]
+
+ return deformed
diff --git a/legacy/Train_and_Inference/utils/coordinate.py b/legacy/Train_and_Inference/utils/coordinate.py
new file mode 100644
index 0000000..fc77f31
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/coordinate.py
@@ -0,0 +1,131 @@
+import numbers
+
+class Coordinate(tuple):
+ '''A ``tuple`` of integers.
+ Allows the following element-wise operators: addition, subtraction,
+ multiplication, division, absolute value, and negation. This allows to
+ perform simple arithmetics with coordinates, e.g.::
+ shape = Coordinate((2, 3, 4))
+ voxel_size = Coordinate((10, 5, 1))
+ size = shape*voxel_size # == Coordinate((20, 15, 4))
+ '''
+ def __new__(cls, array_like):
+ return super(Coordinate, cls).__new__(
+ cls,
+ [
+ int(x)
+ if x is not None
+ else None
+ for x in array_like])
+
+ def dims(self):
+ return len(self)
+
+ def __neg__(self):
+ return Coordinate(
+ -a
+ if a is not None
+ else None
+ for a in self)
+
+ def __abs__(self):
+ return Coordinate(
+ abs(a)
+ if a is not None
+ else None
+ for a in self)
+
+ def __add__(self, other):
+ assert isinstance(
+ other, tuple), "can only add Coordinate or tuples to Coordinate"
+ assert self.dims() == len(other), "can only add Coordinate of equal dimensions"
+ return Coordinate(
+ a+b
+ if a is not None and b is not None
+ else None
+ for a, b in zip(self, other))
+
+ def __sub__(self, other):
+ assert isinstance(
+ other, tuple), "can only subtract Coordinate or tuples to Coordinate"
+ assert self.dims() == len(other), "can only subtract Coordinate of equal dimensions"
+ return Coordinate(
+ a-b
+ if a is not None and b is not None
+ else None
+ for a, b in zip(self, other))
+
+ def __mul__(self, other):
+ if isinstance(other, tuple):
+ assert self.dims() == len(other), "can only multiply Coordinate of equal dimensions"
+ return Coordinate(
+ a*b
+ if a is not None and b is not None
+ else None
+ for a, b in zip(self, other))
+
+ elif isinstance(other, numbers.Number):
+ return Coordinate(
+ a*other
+ if a is not None
+ else None
+ for a in self)
+ else:
+ raise TypeError(
+ "multiplication of Coordinate with type %s not supported" % type(other))
+
+ def __div__(self, other):
+ if isinstance(other, tuple):
+ assert self.dims() == len(other), "can only divide Coordinate of equal dimensions"
+ return Coordinate(
+ a/b
+ if a is not None and b is not None
+ else None
+ for a, b in zip(self, other))
+ elif isinstance(other, numbers.Number):
+ return Coordinate(
+ a/other
+ if a is not None
+ else None
+ for a in self)
+ else:
+ raise TypeError(
+ "division of Coordinate with type %s not supported" % type(other))
+
+ def __truediv__(self, other):
+ if isinstance(other, tuple):
+ assert self.dims() == len(other), "can only divide Coordinate of equal dimensions"
+ return Coordinate(
+ a/b
+ if a is not None and b is not None
+ else None
+ for a, b in zip(self, other))
+ elif isinstance(other, numbers.Number):
+ return Coordinate(
+ a/other
+ if a is not None
+ else None
+ for a in self)
+ else:
+ raise TypeError(
+ "division of Coordinate with type %s not supported" % type(other))
+
+ def __floordiv__(self, other):
+ if isinstance(other, tuple):
+ assert self.dims() == len(other), "can only divide Coordinate of equal dimensions"
+ return Coordinate(
+ a//b
+ if a is not None and b is not None
+ else None
+ for a, b in zip(self, other))
+ elif isinstance(other, numbers.Number):
+ return Coordinate(
+ a//other
+ if a is not None
+ else None
+ for a in self)
+ else:
+ raise TypeError(
+ "division of Coordinate with type %s not supported" % type(other))
+
+
diff --git a/legacy/Train_and_Inference/utils/encoder_dict.py b/legacy/Train_and_Inference/utils/encoder_dict.py
new file mode 100644
index 0000000..3aedf4e
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/encoder_dict.py
@@ -0,0 +1,159 @@
+import torch
+
+ENCODER_DICT = [
+ 'embed_in.0.weight',
+ 'embed_in.0.bias',
+ 'conv0.block1.0.weight',
+ 'conv0.block1.1.weight',
+ 'conv0.block1.1.bias',
+ 'conv0.block2.0.weight',
+ 'conv0.block2.1.weight',
+ 'conv0.block2.1.bias',
+ 'conv0.block2.3.weight',
+ 'conv0.block3.weight',
+ 'conv0.block3.bias',
+ 'conv1.block1.0.weight',
+ 'conv1.block1.1.weight',
+ 'conv1.block1.1.bias',
+ 'conv1.block2.0.weight',
+ 'conv1.block2.1.weight',
+ 'conv1.block2.1.bias',
+ 'conv1.block2.3.weight',
+ 'conv1.block3.weight',
+ 'conv1.block3.bias',
+ 'conv2.block1.0.weight',
+ 'conv2.block1.1.weight',
+ 'conv2.block1.1.bias',
+ 'conv2.block2.0.weight',
+ 'conv2.block2.1.weight',
+ 'conv2.block2.1.bias',
+ 'conv2.block2.3.weight',
+ 'conv2.block3.weight',
+ 'conv2.block3.bias',
+ 'conv3.block1.0.weight',
+ 'conv3.block1.1.weight',
+ 'conv3.block1.1.bias',
+ 'conv3.block2.0.weight',
+ 'conv3.block2.1.weight',
+ 'conv3.block2.1.bias',
+ 'conv3.block2.3.weight',
+ 'conv3.block3.weight',
+ 'conv3.block3.bias',
+ 'center.block1.0.weight',
+ 'center.block1.1.weight',
+ 'center.block1.1.bias',
+ 'center.block2.0.weight',
+ 'center.block2.1.weight',
+ 'center.block2.1.bias',
+ 'center.block2.3.weight',
+ 'center.block3.weight',
+ 'center.block3.bias'
+]
+
+ENCODER_DICT2 = [
+ 'embed_in',
+ 'conv0',
+ 'conv1',
+ 'conv2',
+ 'conv3',
+ 'center'
+]
+
+ENCODER_DECODER_DICT2 = [
+ 'embed_in',
+ 'conv0',
+ 'conv1',
+ 'conv2',
+ 'conv3',
+ 'center',
+ 'up0',
+ 'cat0',
+ 'conv4',
+ 'up1',
+ 'cat1',
+ 'conv5',
+ 'up2',
+ 'cat2',
+ 'conv6',
+ 'up3',
+ 'cat3',
+ 'conv7',
+ 'embed_out'
+]
+
+def freeze_layers(model, if_skip='False'):
+ print('Freeze encoder!')
+ for param in model.embed_in.parameters():
+ param.requires_grad = False
+ for param in model.conv0.parameters():
+ param.requires_grad = False
+ for param in model.conv1.parameters():
+ param.requires_grad = False
+ for param in model.conv2.parameters():
+ param.requires_grad = False
+ for param in model.conv3.parameters():
+ param.requires_grad = False
+ for param in model.center.parameters():
+ param.requires_grad = False
+ if if_skip == 'True':
+ print('Freeze encoder and decoder!')
+ for param in model.up0.parameters():
+ param.requires_grad = False
+ for param in model.cat0.parameters():
+ param.requires_grad = False
+ for param in model.conv4.parameters():
+ param.requires_grad = False
+ for param in model.up1.parameters():
+ param.requires_grad = False
+ for param in model.cat1.parameters():
+ param.requires_grad = False
+ for param in model.conv5.parameters():
+ param.requires_grad = False
+ for param in model.up2.parameters():
+ param.requires_grad = False
+ for param in model.cat2.parameters():
+ param.requires_grad = False
+ for param in model.conv6.parameters():
+ param.requires_grad = False
+ for param in model.up3.parameters():
+ param.requires_grad = False
+ for param in model.cat3.parameters():
+ param.requires_grad = False
+ for param in model.conv7.parameters():
+ param.requires_grad = False
+ for param in model.embed_out.parameters():
+ param.requires_grad = False
+ for param in model.out_put.parameters():
+ param.requires_grad = False
+ return model
+
+def difflr_optimizer(model, lr_base=1e-4, lr_encoder=1e-5, if_skip='False'):
+ print('Adjust the LR of encoder!')
+ encoder_layers_param = []
+ encoder_layers_param += list(map(id, model.embed_in.parameters()))
+ encoder_layers_param += list(map(id, model.conv0.parameters()))
+ encoder_layers_param += list(map(id, model.conv1.parameters()))
+ encoder_layers_param += list(map(id, model.conv2.parameters()))
+ encoder_layers_param += list(map(id, model.conv3.parameters()))
+ encoder_layers_param += list(map(id, model.center.parameters()))
+ if if_skip == 'True':
+ print('Adjust the LR of encoder and decoder!')
+ encoder_layers_param += list(map(id, model.up0.parameters()))
+ encoder_layers_param += list(map(id, model.cat0.parameters()))
+ encoder_layers_param += list(map(id, model.conv4.parameters()))
+ encoder_layers_param += list(map(id, model.up1.parameters()))
+ encoder_layers_param += list(map(id, model.cat1.parameters()))
+ encoder_layers_param += list(map(id, model.conv5.parameters()))
+ encoder_layers_param += list(map(id, model.up2.parameters()))
+ encoder_layers_param += list(map(id, model.cat2.parameters()))
+ encoder_layers_param += list(map(id, model.conv6.parameters()))
+ encoder_layers_param += list(map(id, model.up3.parameters()))
+ encoder_layers_param += list(map(id, model.cat3.parameters()))
+ encoder_layers_param += list(map(id, model.conv7.parameters()))
+ encoder_layers_param += list(map(id, model.embed_out.parameters()))
+ encoder_param = filter(lambda p: id(p) in encoder_layers_param, model.parameters())
+ decoder_param = filter(lambda p: id(p) not in encoder_layers_param, model.parameters())
+ optimizer = torch.optim.Adam([{'params': encoder_param, 'lr': lr_encoder},
+ {'params': decoder_param}],
+ lr=lr_base, betas=(0.9, 0.999), eps=0.01, weight_decay=1e-6, amsgrad=True)
+ return optimizer
diff --git a/legacy/Train_and_Inference/utils/flow_display.py b/legacy/Train_and_Inference/utils/flow_display.py
new file mode 100644
index 0000000..3ae85ee
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/flow_display.py
@@ -0,0 +1,180 @@
+import numpy as np
+import matplotlib.pyplot as plt
+
+def make_color_wheel():
+ """
+ Generate color wheel according Middlebury color code
+ :return: Color wheel
+ """
+ RY = 15
+ YG = 6
+ GC = 4
+ CB = 11
+ BM = 13
+ MR = 6
+
+ ncols = RY + YG + GC + CB + BM + MR
+
+ colorwheel = np.zeros([ncols, 3])
+
+ col = 0
+
+ # RY
+ colorwheel[0:RY, 0] = 255
+ colorwheel[0:RY, 1] = np.transpose(np.floor(255*np.arange(0, RY) / RY))
+ col += RY
+
+ # YG
+ colorwheel[col:col+YG, 0] = 255 - np.transpose(np.floor(255*np.arange(0, YG) / YG))
+ colorwheel[col:col+YG, 1] = 255
+ col += YG
+
+ # GC
+ colorwheel[col:col+GC, 1] = 255
+ colorwheel[col:col+GC, 2] = np.transpose(np.floor(255*np.arange(0, GC) / GC))
+ col += GC
+
+ # CB
+ colorwheel[col:col+CB, 1] = 255 - np.transpose(np.floor(255*np.arange(0, CB) / CB))
+ colorwheel[col:col+CB, 2] = 255
+ col += CB
+
+ # BM
+ colorwheel[col:col+BM, 2] = 255
+ colorwheel[col:col+BM, 0] = np.transpose(np.floor(255*np.arange(0, BM) / BM))
+ col += + BM
+
+ # MR
+ colorwheel[col:col+MR, 2] = 255 - np.transpose(np.floor(255 * np.arange(0, MR) / MR))
+ colorwheel[col:col+MR, 0] = 255
+
+ return colorwheel
+
+def compute_color(u, v):
+ """
+ compute optical flow color map
+ :param u: optical flow horizontal map
+ :param v: optical flow vertical map
+ :return: optical flow in color code
+ """
+ [h, w] = u.shape
+ img = np.zeros([h, w, 3])
+ nanIdx = np.isnan(u) | np.isnan(v)
+ u[nanIdx] = 0
+ v[nanIdx] = 0
+
+ colorwheel = make_color_wheel()
+ ncols = np.size(colorwheel, 0)
+
+ rad = np.sqrt(u**2+v**2)
+
+ a = np.arctan2(-v, -u) / np.pi
+
+ fk = (a+1) / 2 * (ncols - 1) + 1
+
+ k0 = np.floor(fk).astype(int)
+
+ k1 = k0 + 1
+ k1[k1 == ncols+1] = 1
+ f = fk - k0
+
+ for i in range(0, np.size(colorwheel,1)):
+ tmp = colorwheel[:, i]
+ col0 = tmp[k0-1] / 255
+ col1 = tmp[k1-1] / 255
+ col = (1-f) * col0 + f * col1
+
+ idx = rad <= 1
+ col[idx] = 1-rad[idx]*(1-col[idx])
+ notidx = np.logical_not(idx)
+
+ col[notidx] *= 0.75
+ img[:, :, i] = np.uint8(np.floor(255 * col*(1-nanIdx)))
+
+ return img
+
+def flow_to_image(flow):
+ """
+ Convert flow into middlebury color code image
+ :param flow: optical flow map
+ :return: optical flow image in middlebury color
+ """
+ u = flow[:, :, 0]
+ v = flow[:, :, 1]
+
+ maxu = -999.
+ maxv = -999.
+ minu = 999.
+ minv = 999.
+ UNKNOWN_FLOW_THRESH = 1e7
+ SMALLFLOW = 0.0
+ LARGEFLOW = 1e8
+
+ idxUnknow = (abs(u) > UNKNOWN_FLOW_THRESH) | (abs(v) > UNKNOWN_FLOW_THRESH)
+ u[idxUnknow] = 0
+ v[idxUnknow] = 0
+
+ maxu = max(maxu, np.max(u))
+ minu = min(minu, np.min(u))
+
+ maxv = max(maxv, np.max(v))
+ minv = min(minv, np.min(v))
+
+ rad = np.sqrt(u ** 2 + v ** 2)
+ maxrad = max(-1, np.max(rad))
+
+ u = u/(maxrad + np.finfo(float).eps)
+ v = v/(maxrad + np.finfo(float).eps)
+
+ img = compute_color(u, v)
+
+ idx = np.repeat(idxUnknow[:, :, np.newaxis], 3, axis=2)
+ img[idx] = 0
+
+ return np.uint8(img)
+
+def dense_flow(flow):
+ flow_img = flow_to_image(flow)
+ return flow_img
+ # plt.figure()
+ # plt.imshow(flow_img)
+ # plt.axis('off')
+ # plt.show()
+
+def sparse_flow(flow, X=None, Y=None, stride=1):
+ flow = flow.copy()
+ flow[:,:,0] = -flow[:,:,0]
+ if X is None:
+ height, width, _ = flow.shape
+ xx = np.arange(0,height,stride)
+ yy = np.arange(0,width,stride)
+ X, Y= np.meshgrid(xx,yy)
+ X = X.flatten()
+ Y = Y.flatten()
+
+ # sample
+ sample_0 = flow[:, :, 0][xx]
+ sample_0 = sample_0.T
+ sample_x = sample_0[yy]
+ sample_x = sample_x.T
+ sample_1 = flow[:, :, 1][xx]
+ sample_1 = sample_1.T
+ sample_y = sample_1[yy]
+ sample_y = sample_y.T
+
+ sample_x = sample_x[:,:,np.newaxis]
+ sample_y = sample_y[:,:,np.newaxis]
+ new_flow = np.concatenate([sample_x, sample_y], axis=2)
+ flow_x = new_flow[:, :, 0].flatten()
+ flow_y = new_flow[:, :, 1].flatten()
+
+ # display
+ ax = plt.gca()
+ ax.xaxis.set_ticks_position('top')
+ ax.invert_yaxis()
+ # plt.quiver(X,Y, flow_x, flow_y, angles="xy", color="#666666")
+ ax.quiver(X,Y, flow_x, flow_y, color="#666666")
+ ax.grid()
+ # ax.legend()
+ plt.draw()
+ plt.show()
\ No newline at end of file
diff --git a/legacy/Train_and_Inference/utils/flow_synthesis.py b/legacy/Train_and_Inference/utils/flow_synthesis.py
new file mode 100644
index 0000000..6396354
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/flow_synthesis.py
@@ -0,0 +1,161 @@
+import os
+import math
+import random
+import numpy as np
+from PIL import Image
+import matplotlib.pyplot as plt
+from utils.flow_display import dense_flow, sparse_flow
+from utils.image_warp import image_warp
+
+mina = 0.000000001
+def gen_line(p1, p2):
+ denominator = p2[1] - p1[1]
+ if denominator == 0:
+ denominator = mina
+ k = (p2[0] - p1[0]) / denominator
+ b = p1[0] - (k * p1[1])
+ return k, b
+
+def func_line(x, k, b):
+ y = k * x + b
+ return y
+
+def gen_flow(height, width, k, b, line_width=5, fold_width=10, dis_k=0.1):
+ grid_x = np.tile(np.expand_dims(np.arange(width), 0), [height, 1])
+ grid_y = np.tile(np.expand_dims(np.arange(height), 1), [1, width])
+ pos_x = grid_x.flatten()
+ pos_y = grid_y.flatten()
+ dis = (k * pos_x - pos_y + b) / (math.sqrt(k ** 2 + 1))
+ # dis = 1 / dis
+ dis = dis.reshape((height, width))
+ sign = np.zeros_like(dis)
+ mask = np.zeros_like(dis)
+ sign[dis > 0] = 1
+ sign[dis < 0] = -1
+
+ dis_abs = np.abs(dis)
+ mask[dis_abs <= line_width] = 0
+ mask[dis_abs > line_width] = 1
+
+ # dis = dis * 10
+ # max_dis = np.max(dis)
+ # min_dis = np.min(dis)
+ # print(max_dis, min_dis)
+ # dis = (max_dis - dis) * sign + (min_dis - dis) * (1 - sign)
+
+ mask_dis = np.ones_like(dis)
+ mask_dis2 = np.ones_like(dis)
+ dis_width = fold_width - line_width
+ mask_dis[dis_abs < line_width] = 0
+ mask_dis[dis_abs >= line_width] = 1
+ mask_dis2[dis_abs < fold_width] = 0
+ mask_dis2[dis_abs >= fold_width] = 1
+
+ # dis_abs[dis_abs >= dis_width] = fold_width
+ dis_abs_s = np.zeros_like(dis_abs)
+ dis_abs_s2 = np.zeros_like(dis_abs)
+ dis_k = -dis_k
+ dis_b = dis_width - dis_k * line_width
+ dis_abs_s = dis_k * dis_abs + dis_b
+ dis_abs_s[dis_abs_s < 0] = 0
+ dis_abs_s2 = dis_abs_s * mask_dis2 + dis_abs * (1 - mask_dis2)
+ dis_abs_s = dis_abs_s * mask_dis + dis_abs * (1 - mask_dis)
+
+ dis = dis_abs_s * sign
+ dis2 = dis_abs_s2 * (-sign)
+
+ if k == 0:
+ k_T = 1 / mina
+ else:
+ k_T = 1 / k
+
+ angle = angle = math.atan(k_T)
+ sin_p = math.sin(angle)
+ cos_p = math.cos(angle)
+
+ flow = np.zeros((height, width, 2), dtype=np.float32)
+ flow2 = np.zeros((height, width, 2), dtype=np.float32)
+ if k > 0:
+ flow[:, :, 0] = (dis * cos_p)
+ flow[:, :, 1] = -(dis * sin_p)
+ flow2[:, :, 0] = (dis2 * cos_p)
+ flow2[:, :, 1] = -(dis2 * sin_p)
+ else:
+ flow[:, :, 0] = -(dis * cos_p)
+ flow[:, :, 1] = (dis * sin_p)
+ flow2[:, :, 0] = -(dis2 * cos_p)
+ flow2[:, :, 1] = (dis2 * sin_p)
+
+ # print(np.max(flow), np.min(flow))
+ return flow, flow2, mask
+
+if __name__ == "__main__":
+ height = 256
+ width = 256
+ # line_width = 10
+ # fold_width = 20
+ for kkk in range(100):
+ line_width = random.randint(5, 20)
+ fold_width = random.randint(line_width+1, 80)
+
+ # two end points
+ # 1 --> top line (0, x)
+ # 2 --> right line (x, width)
+ # 3 --> bottom line (height, x)
+ # 4 --> left line (x, 0)
+ k1 = random.randint(1, 4)
+ k2 = random.randint(1, 4)
+ while k1 == k2:
+ k2 = random.randint(1, 4)
+
+ if k1 == 1:
+ x = random.randint(1, width-1)
+ p1 = [0, x]
+ elif k1 == 2:
+ x = random.randint(1, height-1)
+ p1 = [x, width]
+ elif k1 == 3:
+ x = random.randint(1, width-1)
+ p1 = [height, x]
+ else:
+ x = random.randint(1, height-1)
+ p1 = [x, 0]
+
+ if k2 == 1:
+ x = random.randint(1, width-1)
+ p2 = [0, x]
+ elif k2 == 2:
+ x = random.randint(1, height-1)
+ p2 = [x, width]
+ elif k2 == 3:
+ x = random.randint(1, width-1)
+ p2 = [height, x]
+ else:
+ x = random.randint(1, height-1)
+ p2 = [x, 0]
+
+ # p1 = [0, 128]
+ # p2 = [128, 256]
+
+ # dis_k = random.uniform(0.001, 0.1)
+ dis_k = random.uniform(0.00001, 0.1)
+ k, b = gen_line(p1, p2)
+
+ flow, flow2, mask = gen_flow(height, width, k, b, line_width, fold_width, dis_k)
+ # flow = flow * 10
+ # print(flow[:10,-10:,0])
+ # print(flow[:10,-10:,1])
+ flow_show1 = dense_flow(flow)
+ flow_show2 = dense_flow(flow2)
+ flow_show = np.concatenate([flow_show1, flow_show2], axis=1)
+ Image.fromarray(flow_show).save('./temp/flow_'+str(kkk).zfill(4)+'.png')
+ # sparse_flow(flow2, stride=10)
+
+
+ # img = np.asarray(Image.open('./0000.png'))
+ # img = img[:256, :256]
+
+ # deformed = image_warp(img, flow, mode='bilinear') # nearest or bilinear
+ # deformed = (deformed * mask).astype(np.uint8)
+ # Image.fromarray(img).save('./deformed1.png')
+ # Image.fromarray(deformed).save('./deformed2.png')
diff --git a/legacy/Train_and_Inference/utils/gen_pseudo.py b/legacy/Train_and_Inference/utils/gen_pseudo.py
new file mode 100644
index 0000000..25b15a9
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/gen_pseudo.py
@@ -0,0 +1,74 @@
+import math
+import torch
+import numpy as np
+from yaml.events import NodeEvent
+
+class GenPseudo(object):
+ def __init__(self, mode='threshold',
+ threshold=0.99,
+ proportion=0.20):
+ super(GenPseudo, self).__init__()
+ self.mode = mode
+ self.threshold = threshold
+ self.proportion = proportion
+
+ def __call__(self, inputs):
+ return self.forward(inputs)
+
+ def forward(self, inputs):
+ inputs = inputs.detach().clone()
+ mask = torch.zeros_like(inputs)
+ if self.mode == 'threshold':
+ inputs[inputs > self.threshold] = 1
+ mask[inputs == 1] = 1
+ inputs[inputs < (1-self.threshold)] = 0
+ mask[inputs == 0] = 1
+ return inputs, mask
+ else:
+ num_classes = 2 # binary classification
+ pseudo_lb = []
+ masks = []
+ batch_size = inputs.shape[0]
+ for k in range(batch_size):
+ output_affs = inputs[k]
+ output_affs_0 = 1 - output_affs.clone()
+ output_affs_1 = output_affs.clone()
+ output_affs_all = torch.stack([output_affs_0, output_affs_1], dim=0)
+ probmap_max, pred_label = torch.max(output_affs_all, dim=0)
+ for idx_cls in range(num_classes):
+ out_div_all = []
+ for i in range(3):
+ pred_label_temp = pred_label[i]
+ probmap_max_temp = probmap_max[i]
+ probmap_max_cls_temp = probmap_max_temp[pred_label_temp == idx_cls]
+ if len(probmap_max_cls_temp) > 0:
+ # probmap_max_cls_temp = probmap_max_cls_temp.view(probmap_max_cls_temp.size(0), -1)
+ probmap_max_cls_temp = probmap_max_cls_temp[0:len(probmap_max_cls_temp)]
+ probmap_max_cls_temp, _ = torch.sort(probmap_max_cls_temp, descending=True)
+ len_cls = len(probmap_max_cls_temp)
+ thresh_len = int(math.floor(len_cls * self.proportion))
+ thresh_temp = probmap_max_cls_temp[thresh_len - 1]
+ out_div = torch.div(output_affs_all[idx_cls, i], thresh_temp)
+ else:
+ out_div = output_affs_all[idx_cls, i]
+ out_div_all.append(out_div)
+ out_div_all = torch.stack(out_div_all, dim=0)
+ output_affs_all[idx_cls] = out_div_all
+
+ rw_probmap_max, pseudo_label = torch.max(output_affs_all, dim=0)
+ mask = torch.zeros_like(rw_probmap_max)
+ mask[rw_probmap_max>=1] = 1
+ pseudo_lb.append(pseudo_label)
+ masks.append(mask)
+ pseudo_lb = torch.stack(pseudo_lb, dim=0)
+ masks = torch.stack(masks, dim=0)
+
+ return pseudo_lb, masks
+
+
+if __name__ == "__main__":
+ gen_pseudo = GenPseudo(mode='prop')
+ pred = np.random.random((2,3,18,160,160)).astype(np.float32)
+ pred = torch.tensor(pred).to('cuda:0')
+
+ pseudo_lb, masks = gen_pseudo(pred)
diff --git a/legacy/Train_and_Inference/utils/image_warp.py b/legacy/Train_and_Inference/utils/image_warp.py
new file mode 100644
index 0000000..1bc473b
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/image_warp.py
@@ -0,0 +1,112 @@
+import numpy as np
+
+def image_warp(im, flow, mode='bilinear'):
+ """Performs a backward warp of an image using the predicted flow.
+ numpy version
+
+ Args:
+ im: input image. ndim=2, 3 or 4, [[num_batch], height, width, [channels]]. num_batch and channels are optional, default is 1.
+ flow: flow vectors. ndim=3 or 4, [[num_batch], height, width, 2]. num_batch is optional
+ mode: interpolation mode. 'nearest' or 'bilinear'
+ Returns:
+ warped: transformed image of the same shape as the input image.
+ """
+ # assert im.ndim == flow.ndim, 'The dimension of im and flow must be equal '
+ flag = 4
+ if im.ndim == 2:
+ height, width = im.shape
+ num_batch = 1
+ channels = 1
+ im = im[np.newaxis, :, :, np.newaxis]
+ flow = flow[np.newaxis, :, :]
+ flag = 2
+ elif im.ndim == 3:
+ height, width, channels = im.shape
+ num_batch = 1
+ im = im[np.newaxis, :, :]
+ flow = flow[np.newaxis, :, :]
+ flag = 3
+ elif im.ndim == 4:
+ num_batch, height, width, channels = im.shape
+ flag = 4
+ else:
+ raise AttributeError('The dimension of im must be 2, 3 or 4')
+
+ max_x = width - 1
+ max_y = height - 1
+ zero = 0
+
+ # We have to flatten our tensors to vectorize the interpolation
+ im_flat = np.reshape(im, [-1, channels])
+ flow_flat = np.reshape(flow, [-1, 2])
+
+ # Floor the flow, as the final indices are integers
+ flow_floor = np.floor(flow_flat).astype(np.int32)
+
+ # Construct base indices which are displaced with the flow
+ pos_x = np.tile(np.arange(width), [height * num_batch])
+ grid_y = np.tile(np.expand_dims(np.arange(height), 1), [1, width])
+ pos_y = np.tile(np.reshape(grid_y, [-1]), [num_batch])
+
+ x = flow_floor[:, 0]
+ y = flow_floor[:, 1]
+
+ x0 = pos_x + x
+ y0 = pos_y + y
+
+ x0 = np.clip(x0, zero, max_x)
+ y0 = np.clip(y0, zero, max_y)
+
+ dim1 = width * height
+ batch_offsets = np.arange(num_batch) * dim1
+ base_grid = np.tile(np.expand_dims(batch_offsets, 1), [1, dim1])
+ base = np.reshape(base_grid, [-1])
+
+ base_y0 = base + y0 * width
+
+ if mode == 'nearest':
+ idx_a = base_y0 + x0
+ warped_flat = im_flat[idx_a]
+ elif mode == 'bilinear':
+ # The fractional part is used to control the bilinear interpolation.
+ bilinear_weights = flow_flat - np.floor(flow_flat)
+
+ xw = bilinear_weights[:, 0]
+ yw = bilinear_weights[:, 1]
+
+ # Compute interpolation weights for 4 adjacent pixels
+ # expand to num_batch * height * width x 1 for broadcasting in add_n below
+ wa = np.expand_dims((1 - xw) * (1 - yw), 1) # top left pixel
+ wb = np.expand_dims((1 - xw) * yw, 1) # bottom left pixel
+ wc = np.expand_dims(xw * (1 - yw), 1) # top right pixel
+ wd = np.expand_dims(xw * yw, 1) # bottom right pixel
+
+ x1 = x0 + 1
+ y1 = y0 + 1
+
+ x1 = np.clip(x1, zero, max_x)
+ y1 = np.clip(y1, zero, max_y)
+
+ base_y1 = base + y1 * width
+ idx_a = base_y0 + x0
+ idx_b = base_y1 + x0
+ idx_c = base_y0 + x1
+ idx_d = base_y1 + x1
+
+ Ia = im_flat[idx_a]
+ Ib = im_flat[idx_b]
+ Ic = im_flat[idx_c]
+ Id = im_flat[idx_d]
+
+ warped_flat = wa * Ia + wb * Ib + wc * Ic + wd * Id
+ warped = np.reshape(warped_flat, [num_batch, height, width, channels])
+
+ if flag == 2:
+ warped = np.squeeze(warped)
+ elif flag == 3:
+ warped = np.squeeze(warped, axis=0)
+ else:
+ pass
+ warped = warped.astype(np.uint8)
+
+ return warped
\ No newline at end of file
diff --git a/legacy/Train_and_Inference/utils/malis_loss.py b/legacy/Train_and_Inference/utils/malis_loss.py
new file mode 100644
index 0000000..6a0eba0
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/malis_loss.py
@@ -0,0 +1,14 @@
+import numpy as np
+from em_segLib.seg_malis import malis_init, malis_loss_weights_both
+from em_segLib.seg_util import mknhood3d
+
+def malis_loss(output_affs, test_label, seg):
+ seg = seg.astype(np.uint64)
+ conn_dims = np.array(output_affs.shape).astype(np.uint64)
+ nhood_dims = np.array((3,3),dtype=np.uint64)
+ nhood_data = mknhood3d(1).astype(np.int32).flatten()
+ pre_ve, pre_prodDims, pre_nHood = malis_init(conn_dims, nhood_data, nhood_dims)
+ weight = malis_loss_weights_both(seg.flatten(), conn_dims, nhood_data, nhood_dims, pre_ve,
+ pre_prodDims, pre_nHood, output_affs.flatten(), test_label.flatten(), 0.5).reshape(conn_dims)
+ malis = np.sum(weight * (output_affs - test_label) ** 2)
+ return malis
diff --git a/legacy/Train_and_Inference/utils/optim_weight_ema.py b/legacy/Train_and_Inference/utils/optim_weight_ema.py
new file mode 100644
index 0000000..bf2b1c4
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/optim_weight_ema.py
@@ -0,0 +1,25 @@
+import torch
+
+
+class EMAWeightOptimizer (object):
+ def __init__(self, target_net, source_net, ema_alpha):
+ self.target_net = target_net
+ self.source_net = source_net
+ self.ema_alpha = ema_alpha
+ self.target_params = [p for p in target_net.state_dict().values() if p.dtype == torch.float]
+ self.source_params = [p for p in source_net.state_dict().values() if p.dtype == torch.float]
+
+ for tgt_p, src_p in zip(self.target_params, self.source_params):
+ tgt_p[...] = src_p[...]
+
+ target_keys = set(target_net.state_dict().keys())
+ source_keys = set(source_net.state_dict().keys())
+ if target_keys != source_keys:
+ raise ValueError('Source and target networks do not have the same state dict keys; do they have different architectures?')
+
+
+ def step(self):
+ one_minus_alpha = 1.0 - self.ema_alpha
+ for tgt_p, src_p in zip(self.target_params, self.source_params):
+ tgt_p.mul_(self.ema_alpha)
+ tgt_p.add_(src_p * one_minus_alpha)
diff --git a/legacy/Train_and_Inference/utils/post_func.py b/legacy/Train_and_Inference/utils/post_func.py
new file mode 100644
index 0000000..372ab1a
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/post_func.py
@@ -0,0 +1,213 @@
+'''
+Descripttion:
+version: 0.0
+Author: Wei Huang
+Date: 2021-11-01 16:24:30
+'''
+import mahotas
+import numpy as np
+
+from scipy import ndimage
+import elf.segmentation.multicut as mc
+import elf.segmentation.features as feats
+import elf.segmentation.watershed as ws
+from scipy.ndimage.morphology import distance_transform_edt
+from scipy.ndimage.filters import gaussian_filter, maximum_filter
+
+# reduce the labeling
+def getSegType(mid):
+ m_type = np.uint64
+ if mid<2**8:
+ m_type = np.uint8
+ elif mid<2**16:
+ m_type = np.uint16
+ elif mid<2**32:
+ m_type = np.uint32
+ return m_type
+
+def relabel(seg, do_type=False):
+ # get the unique labels
+ uid = np.unique(seg)
+ # ignore all-background samples
+ if len(uid)==1 and uid[0] == 0:
+ return seg
+
+ uid = uid[uid > 0]
+ mid = int(uid.max()) + 1 # get the maximum label for the segment
+
+ # create an array from original segment id to reduced id
+ m_type = seg.dtype
+ if do_type:
+ m_type = getSegType(mid)
+ mapping = np.zeros(mid, dtype=m_type)
+ mapping[uid] = np.arange(1, len(uid) + 1, dtype=m_type)
+ return mapping[seg]
+
+def randomlabel(segmentation):
+ segmentation = segmentation.astype(np.uint32)
+ uid = np.unique(segmentation)
+ mid = int(uid.max()) + 1
+ mapping = np.zeros(mid, dtype=segmentation.dtype)
+ mapping[uid] = np.random.choice(len(uid), len(uid), replace=False).astype(segmentation.dtype)#(len(uid), dtype=segmentation.dtype)
+ out = mapping[segmentation]
+ out[segmentation==0] = 0
+ return out
+
+def mc_baseline(affs, fragments=None):
+ affs = 1 - affs
+ boundary_input = np.maximum(affs[1], affs[2])
+ if fragments is None:
+ fragments = np.zeros_like(boundary_input, dtype='uint64')
+ offset = 0
+ for z in range(fragments.shape[0]):
+ wsz, max_id = ws.distance_transform_watershed(boundary_input[z], threshold=.25, sigma_seeds=2.)
+ wsz += offset
+ offset += max_id
+ fragments[z] = wsz
+ rag = feats.compute_rag(fragments)
+ offsets = [[-1, 0, 0], [0, -1, 0], [0, 0, -1]]
+ costs = feats.compute_affinity_features(rag, affs, offsets)[:, 0]
+ edge_sizes = feats.compute_boundary_mean_and_length(rag, boundary_input)[:, 1]
+ costs = mc.transform_probabilities_to_costs(costs, edge_sizes=edge_sizes)
+ node_labels = mc.multicut_kernighan_lin(rag, costs)
+ segmentation = feats.project_node_labels_to_pixels(rag, node_labels)
+ return segmentation
+
+def watershed(affs, seed_method, use_mahotas_watershed=True):
+ affs_xy = 1.0 - 0.5*(affs[1] + affs[2])
+ depth = affs_xy.shape[0]
+ fragments = np.zeros_like(affs[0]).astype(np.uint64)
+ next_id = 1
+ for z in range(depth):
+ seeds, num_seeds = get_seeds(affs_xy[z], next_id=next_id, method=seed_method)
+ if use_mahotas_watershed:
+ fragments[z] = mahotas.cwatershed(affs_xy[z], seeds)
+ else:
+ fragments[z] = ndimage.watershed_ift((255.0*affs_xy[z]).astype(np.uint8), seeds)
+ next_id += num_seeds
+ return fragments
+
+def get_seeds(boundary, method='grid', next_id=1, seed_distance=10):
+ if method == 'grid':
+ height = boundary.shape[0]
+ width = boundary.shape[1]
+ seed_positions = np.ogrid[0:height:seed_distance, 0:width:seed_distance]
+ num_seeds_y = seed_positions[0].size
+ num_seeds_x = seed_positions[1].size
+ num_seeds = num_seeds_x*num_seeds_y
+ seeds = np.zeros_like(boundary).astype(np.int32)
+ seeds[seed_positions] = np.arange(next_id, next_id + num_seeds).reshape((num_seeds_y,num_seeds_x))
+
+ if method == 'minima':
+ minima = mahotas.regmin(boundary)
+ seeds, num_seeds = mahotas.label(minima)
+ seeds += next_id
+ seeds[seeds==next_id] = 0
+
+ if method == 'maxima_distance':
+ distance = mahotas.distance(boundary<0.5)
+ maxima = mahotas.regmax(distance)
+ seeds, num_seeds = mahotas.label(maxima)
+ seeds += next_id
+ seeds[seeds==next_id] = 0
+
+ return seeds, num_seeds
+
+
+def watershed_lmc(affs):
+ affs = 1 - affs
+ boundary_input = np.maximum(affs[1], affs[2])
+ fragments = np.zeros_like(boundary_input, dtype=np.uint64)
+ offset = 0
+ for z in range(fragments.shape[0]):
+ wsz, max_id = ws.distance_transform_watershed(boundary_input[z], threshold=.25, sigma_seeds=2.)
+ wsz += offset
+ offset += max_id
+ fragments[z] = wsz
+ return fragments, offset
+
+
+def agglomerate_lmc(affs, fragments):
+ affs = 1 - affs
+ boundary_input = np.maximum(affs[1], affs[2])
+ rag = feats.compute_rag(fragments)
+ offsets = [[-1, 0, 0], [0, -1, 0], [0, 0, -1]]
+ costs = feats.compute_affinity_features(rag, affs, offsets)[:, 0]
+ edge_sizes = feats.compute_boundary_mean_and_length(rag, boundary_input)[:, 1]
+ costs = mc.transform_probabilities_to_costs(costs, edge_sizes=edge_sizes)
+ node_labels = mc.multicut_kernighan_lin(rag, costs)
+ segmentation = feats.project_node_labels_to_pixels(rag, node_labels)
+ return segmentation
+
+
+# copy from LSD --> fragments.py
+def watershed_from_affinities(
+ affs,
+ max_affinity_value=1.0,
+ fragments_in_xy=True,
+ return_seeds=False,
+ min_seed_distance=10):
+ '''Extract initial fragments from affinities using a watershed
+ transform. Returns the fragments and the maximal ID in it.
+
+ Returns:
+ (fragments, max_id)
+ or
+ (fragments, max_id, seeds) if return_seeds == True'''
+
+ if fragments_in_xy:
+ mean_affs = 0.5 * (affs[1] + affs[2])
+ depth = mean_affs.shape[0]
+ fragments = np.zeros(mean_affs.shape, dtype=np.uint64)
+ if return_seeds:
+ seeds = np.zeros(mean_affs.shape, dtype=np.uint64)
+ id_offset = 0
+ for z in range(depth):
+ boundary_mask = mean_affs[z] > 0.5 * max_affinity_value
+ boundary_distances = distance_transform_edt(boundary_mask)
+ ret = watershed_from_boundary_distance(
+ boundary_distances,
+ return_seeds=return_seeds,
+ id_offset=id_offset,
+ min_seed_distance=min_seed_distance)
+ fragments[z] = ret[0]
+ if return_seeds:
+ seeds[z] = ret[2]
+ id_offset = ret[1]
+ ret = (fragments, id_offset)
+ if return_seeds:
+ ret += (seeds,)
+ else:
+ boundary_mask = np.mean(affs, axis=0) > 0.5 * max_affinity_value
+ boundary_distances = distance_transform_edt(boundary_mask)
+ ret = watershed_from_boundary_distance(
+ boundary_distances,
+ return_seeds=return_seeds,
+ min_seed_distance=min_seed_distance)
+ fragments = ret[0]
+ return ret
+
+
+def watershed_from_boundary_distance(
+ boundary_distances,
+ return_seeds=False,
+ id_offset=0,
+ min_seed_distance=10):
+ max_filtered = maximum_filter(boundary_distances, min_seed_distance)
+ maxima = max_filtered == boundary_distances
+ seeds, n = mahotas.label(maxima)
+
+ if n == 0:
+ return np.zeros(boundary_distances.shape, dtype=np.uint64), id_offset
+
+ seeds[seeds!=0] += id_offset
+
+ fragments = mahotas.cwatershed(
+ boundary_distances.max() - boundary_distances,
+ seeds)
+
+ ret = (fragments.astype(np.uint64), n + id_offset)
+ if return_seeds:
+ ret = ret + (seeds.astype(np.uint64),)
+
+ return ret
diff --git a/legacy/Train_and_Inference/utils/post_lmc.py b/legacy/Train_and_Inference/utils/post_lmc.py
new file mode 100644
index 0000000..b846709
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/post_lmc.py
@@ -0,0 +1,77 @@
+import numpy as np
+from skimage.metrics import adapted_rand_error as adapted_rand_ref
+from skimage.metrics import variation_of_information as voi_ref
+import elf.segmentation.watershed as ws
+import elf.segmentation.multicut as mc
+import elf.segmentation.features as feats
+import elf.segmentation.watershed as ws
+from elf.segmentation.features import *
+from elf.segmentation.learning import *
+from elf.segmentation.mutex_watershed import mutex_watershed
+from elf.parallel.relabel import relabel_consecutive
+from nifty import tools as ntools
+import nifty.graph.rag as nrag
+import os
+import time
+from tqdm import tqdm
+import numpy as np
+import joblib
+import imageio
+from skimage.metrics import variation_of_information, adapted_rand_error
+from multiprocessing import Pool, Lock
+
+def post_lmc(affs):
+ affs = 1 - affs
+ boundary_input = np.maximum(affs[1], affs[2])
+ watershed = np.zeros_like(boundary_input, dtype='uint64')
+ offset = 0
+ for z in range(watershed.shape[0]):
+ wsz, max_id = ws.distance_transform_watershed(boundary_input[z], threshold=.25, sigma_seeds=2.)
+ wsz += offset
+ offset += max_id
+ watershed[z] = wsz
+ rag = feats.compute_rag(watershed)
+ offsets = [[-1, 0, 0], [0, -1, 0], [0, 0, -1]]
+ costs = feats.compute_affinity_features(rag, affs, offsets)[:, 0]
+ edge_sizes = feats.compute_boundary_mean_and_length(rag, boundary_input)[:, 1]
+ costs = mc.transform_probabilities_to_costs(costs, edge_sizes=edge_sizes, beta=0.25)
+ node_labels = mc.multicut_kernighan_lin(rag, costs)
+ segmentation = feats.project_node_labels_to_pixels(rag, node_labels)
+ return segmentation
+
+def post_lmc_lh(affs, beta):
+ affs = 1 - affs
+ boundary_input = np.maximum(affs[1], affs[2])
+ watershed = np.zeros_like(boundary_input, dtype='uint64')
+ offset = 0
+ for z in range(watershed.shape[0]):
+ wsz, max_id = ws.distance_transform_watershed(boundary_input[z], threshold=.25, sigma_seeds=2.)
+ wsz += offset
+ offset += max_id
+ watershed[z] = wsz
+ rag = feats.compute_rag(watershed)
+ offsets = [[-1, 0, 0], [0, -1, 0], [0, 0, -1]]
+ costs = feats.compute_affinity_features(rag, affs, offsets)[:, 0]
+ edge_sizes = feats.compute_boundary_mean_and_length(rag, boundary_input)[:, 1]
+ costs = mc.transform_probabilities_to_costs(costs, edge_sizes=edge_sizes, beta=beta)
+ node_labels = mc.multicut_kernighan_lin(rag, costs)
+ segmentation = feats.project_node_labels_to_pixels(rag, node_labels)
+ return segmentation
+
+def post_mc_b(boundary_input, beta=0.25):
+ boundary_input = 1 - boundary_input
+ watershed = np.zeros_like(boundary_input, dtype='uint64')
+ offset = 0
+ for z in range(watershed.shape[0]):
+ wsz, max_id = ws.distance_transform_watershed(boundary_input[z], threshold=0.25, sigma_seeds=2.0)
+ wsz += offset
+ offset += max_id
+ watershed[z] = wsz
+ rag = feats.compute_rag(watershed)
+ costs = compute_boundary_features(rag, boundary_input, min_value=0, max_value=1)[:, 0]
+ edge_sizes = feats.compute_boundary_mean_and_length(rag, boundary_input)[:, 1]
+ costs = mc.transform_probabilities_to_costs(costs, edge_sizes=edge_sizes, beta=beta)
+ node_labels = mc.multicut_kernighan_lin(rag, costs)
+ segmentation = feats.project_node_labels_to_pixels(rag, node_labels)
+
+ return segmentation
\ No newline at end of file
diff --git a/legacy/Train_and_Inference/utils/post_waterz.py b/legacy/Train_and_Inference/utils/post_waterz.py
new file mode 100644
index 0000000..439de58
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/post_waterz.py
@@ -0,0 +1,92 @@
+import waterz
+import mahotas
+import numpy as np
+from scipy import ndimage
+
+def randomlabel(segmentation):
+ segmentation = segmentation.astype(np.uint32)
+ uid = np.unique(segmentation)
+ mid = int(uid.max()) + 1
+ mapping = np.zeros(mid, dtype=segmentation.dtype)
+ mapping[uid] = np.random.choice(len(uid), len(uid), replace=False).astype(segmentation.dtype)#(len(uid), dtype=segmentation.dtype)
+ out = mapping[segmentation]
+ out[segmentation==0] = 0
+ return out
+
+def watershed(affs, seed_method, use_mahotas_watershed=True):
+ affs_xy = 1.0 - 0.5*(affs[1] + affs[2])
+ depth = affs_xy.shape[0]
+ fragments = np.zeros_like(affs[0]).astype(np.uint64)
+ next_id = 1
+ for z in range(depth):
+ seeds, num_seeds = get_seeds(affs_xy[z], next_id=next_id, method=seed_method)
+ if use_mahotas_watershed:
+ fragments[z] = mahotas.cwatershed(affs_xy[z], seeds)
+ else:
+ fragments[z] = ndimage.watershed_ift((255.0*affs_xy[z]).astype(np.uint8), seeds)
+ next_id += num_seeds
+ return fragments
+
+def get_seeds(boundary, method='grid', next_id=1, seed_distance=10):
+ if method == 'grid':
+ height = boundary.shape[0]
+ width = boundary.shape[1]
+ seed_positions = np.ogrid[0:height:seed_distance, 0:width:seed_distance]
+ num_seeds_y = seed_positions[0].size
+ num_seeds_x = seed_positions[1].size
+ num_seeds = num_seeds_x*num_seeds_y
+ seeds = np.zeros_like(boundary).astype(np.int32)
+ seeds[seed_positions] = np.arange(next_id, next_id + num_seeds).reshape((num_seeds_y,num_seeds_x))
+
+ if method == 'minima':
+ minima = mahotas.regmin(boundary)
+ seeds, num_seeds = mahotas.label(minima)
+ seeds += next_id
+ seeds[seeds==next_id] = 0
+
+ if method == 'maxima_distance':
+ distance = mahotas.distance(boundary<0.5)
+ maxima = mahotas.regmax(distance)
+ seeds, num_seeds = mahotas.label(maxima)
+ seeds += next_id
+ seeds[seeds==next_id] = 0
+
+ return seeds, num_seeds
+
+def elf_watershed(affs):
+ import elf.segmentation.watershed as ws
+ affs = 1 - affs
+ boundary_input = np.maximum(affs[1], affs[2])
+ fragments = np.zeros_like(boundary_input, dtype='uint64')
+ offset = 0
+ for z in range(fragments.shape[0]):
+ wsz, max_id = ws.distance_transform_watershed(boundary_input[z], threshold=.25, sigma_seeds=2.)
+ wsz += offset
+ offset += max_id
+ fragments[z] = wsz
+ return fragments
+
+def relabel(seg):
+ # get the unique labels
+ uid = np.unique(seg)
+ # ignore all-background samples
+ if len(uid)==1 and uid[0] == 0:
+ return seg
+
+ uid = uid[uid > 0]
+ mid = int(uid.max()) + 1 # get the maximum label for the segment
+
+ # create an array from original segment id to reduced id
+ m_type = seg.dtype
+ mapping = np.zeros(mid, dtype=m_type)
+ mapping[uid] = np.arange(1, len(uid) + 1, dtype=m_type)
+ return mapping[seg]
+
+def post_waterz(affs, thresd=0.5):
+ fragments = watershed(affs, 'maxima_distance')
+ sf = 'OneMinus>'
+ seg = list(waterz.agglomerate(affs, [0.50],
+ fragments=fragments,
+ scoring_function=sf,
+ discretize_queue=256))[0]
+ return seg
diff --git a/legacy/Train_and_Inference/utils/seeds_func.py b/legacy/Train_and_Inference/utils/seeds_func.py
new file mode 100644
index 0000000..edfe4d7
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/seeds_func.py
@@ -0,0 +1,443 @@
+import os
+import numpy as np
+import h5py
+from scipy import ndimage
+import cv2
+import mahotas
+import matplotlib
+matplotlib.use("agg")
+import matplotlib.pyplot as plt
+
+# generate affinity
+def seg_to_affgraph(seg, nhood=np.array([[-1, 0, 0], [0, -1, 0], [0, 0, -1]])):
+ # constructs an affinity graph from a segmentation
+ # assume affinity graph is represented as:
+ # shape = (e, z, y, x)
+ # nhood.shape = (edges, 3)
+ shape = seg.shape
+ nEdge = nhood.shape[0]
+ aff = np.zeros((nEdge,)+shape,dtype=np.int32)
+
+ for e in range(nEdge):
+ aff[e, \
+ max(0,-nhood[e,0]):min(shape[0],shape[0]-nhood[e,0]), \
+ max(0,-nhood[e,1]):min(shape[1],shape[1]-nhood[e,1]), \
+ max(0,-nhood[e,2]):min(shape[2],shape[2]-nhood[e,2])] = \
+ (seg[max(0,-nhood[e,0]):min(shape[0],shape[0]-nhood[e,0]), \
+ max(0,-nhood[e,1]):min(shape[1],shape[1]-nhood[e,1]), \
+ max(0,-nhood[e,2]):min(shape[2],shape[2]-nhood[e,2])] == \
+ seg[max(0,nhood[e,0]):min(shape[0],shape[0]+nhood[e,0]), \
+ max(0,nhood[e,1]):min(shape[1],shape[1]+nhood[e,1]), \
+ max(0,nhood[e,2]):min(shape[2],shape[2]+nhood[e,2])] ) \
+ * ( seg[max(0,-nhood[e,0]):min(shape[0],shape[0]-nhood[e,0]), \
+ max(0,-nhood[e,1]):min(shape[1],shape[1]-nhood[e,1]), \
+ max(0,-nhood[e,2]):min(shape[2],shape[2]-nhood[e,2])] > 0 ) \
+ * ( seg[max(0,nhood[e,0]):min(shape[0],shape[0]+nhood[e,0]), \
+ max(0,nhood[e,1]):min(shape[1],shape[1]+nhood[e,1]), \
+ max(0,nhood[e,2]):min(shape[2],shape[2]+nhood[e,2])] > 0 )
+
+ return aff
+
+
+
+# generate seeds
+def gen_seeds(labels, affs_xy, min_size=10):
+ # remove some neurons whose size is smaller than min_size
+ ids, count = np.unique(labels, return_counts=True)
+ for i, icount in enumerate(count):
+ if icount < min_size:
+ labels[labels == ids[i]] = 0
+
+ boundary = np.ones_like(affs_xy)
+ boundary[1:-1, 1:-1] = affs_xy[1:-1, 1:-1]
+ boundary[boundary != 0] = 1
+
+ distance = mahotas.distance(boundary<0.5)
+ seeds = np.zeros_like(labels)
+ ite = 1
+ for label in np.unique(labels):
+ if label == 0:
+ continue
+ label_mask = labels == label
+ label_mask = label_mask.astype(np.int)
+ temp_dis = np.multiply(distance, label_mask)
+ max_where = np.where(temp_dis == np.max(temp_dis))
+ seeds[max_where[0][0], max_where[1][0]] = ite
+ ite += 1
+ return seeds, boundary
+
+
+def gen_seeds_2(labels, affs_xy, min_size=10):
+ # remove some neurons whose size is smaller than min_size
+ ids, count = np.unique(labels, return_counts=True)
+ for i, icount in enumerate(count):
+ if icount < min_size:
+ labels[labels == ids[i]] = 0
+
+ boundary = np.ones_like(affs_xy)
+ boundary[1:-1, 1:-1] = affs_xy[1:-1, 1:-1]
+ boundary[boundary != 0] = 1
+
+ distance = mahotas.distance(boundary<0.5)
+ seeds = np.zeros_like(labels)
+ # ite = 1
+ for label in np.unique(labels):
+ if label == 0:
+ continue
+ label_mask = labels == label
+ label_mask = label_mask.astype(np.int)
+ temp_dis = np.multiply(distance, label_mask)
+ max_where = np.where(temp_dis == np.max(temp_dis))
+ seeds[max_where[0][0], max_where[1][0]] = label
+ # ite += 1
+ return seeds
+
+
+# erosion labels
+def erosion_labels(gt, steps=1):
+ self_background = 0
+ foreground = np.zeros(shape=gt.shape, dtype=np.bool)
+ for label in np.unique(gt):
+ if label == self_background:
+ continue
+ label_mask = gt==label
+ # Assume that masked out values are the same as the label we are
+ # eroding in this iteration. This ensures that at the boundary to
+ # a masked region the value blob is not shrinking.
+ eroded_label_mask = ndimage.binary_erosion(label_mask, iterations=steps, border_value=1)
+ foreground = np.logical_or(eroded_label_mask, foreground)
+ background = np.logical_not(foreground)
+ gt[background] = self_background
+ return gt
+
+
+# draw fragments
+def draw_fragments(picture, raw=None, alpha=0.3):
+ m,n = picture.shape
+ ids = np.unique(picture)
+ size = len(ids)
+ print("The number of nuerons is %d" % size)
+ color = np.zeros([m, n, 3])
+ idx = np.searchsorted(ids, picture)
+ for i in range(3):
+ color_val = np.random.randint(0, 255, ids.shape)
+ if ids[0] == 0:
+ color_val[0] = 0
+ color[:,:,i] = color_val[idx]
+ color = color / 255
+ if raw is not None:
+ plt.figure()
+ plt.subplots(figsize=(10,10))
+ plt.imshow(raw)
+ plt.imshow(color, alpha=alpha)
+ plt.axis('off')
+ plt.show()
+ else:
+ plt.figure()
+ plt.subplots(figsize=(10,10))
+ plt.imshow(color)
+ plt.axis('off')
+ plt.show()
+
+
+# thresdholding
+def binary_thresholding(img, t=0.5):
+ if np.max(img) > 1.0:
+ img = img / 255.0
+ img[img >= t] = 1
+ img[img < t] = 0
+ return img
+
+
+# draw seeds
+def draw_seeds(raw, seeds):
+ plt.figure(figsize=(10,10))
+ plt.imshow(raw, cmap='gray')
+ seeds_listx, seeds_listy = np.where(seeds != 0)
+ plt.scatter(seeds_listy, seeds_listx, c='r')
+ plt.axis('off')
+ plt.show()
+
+
+def draw_seeds_v2(raw, seeds):
+ plt.figure(figsize=(10,10))
+ plt.imshow(raw, cmap='gray')
+ seeds_listx = seeds[:, 0].astype(np.int)
+ seeds_listy = seeds[:, 1].astype(np.int)
+ plt.scatter(seeds_listy, seeds_listx, c='r')
+ plt.axis('off')
+ plt.show()
+
+
+def draw_box(img, box):
+ img_box = img.copy()
+ if len(img_box.shape) == 2:
+ img_box = img_box[:,:,np.newaxis]
+ img_box = np.concatenate([img_box, img_box, img_box], axis=2)
+ for i in range(1, box.shape[0]):
+ position = box[i]
+ x1 = position[0]
+ y1 = position[1]
+ x2 = x1 + position[2]
+ y2 = y1 + position[3]
+ img_box = cv2.rectangle(img_box, (int(x1), int(y1)), (int(x2), int(y2)), (255, 0, 0), 2)
+ plt.figure(figsize=(10,10))
+ plt.imshow(img_box)
+ plt.axis('off')
+ plt.show()
+
+
+# draw affinity
+def draw_general(img):
+ plt.figure(figsize=(10,10))
+ plt.imshow(img, cmap='gray')
+ plt.axis('off')
+ plt.show()
+
+
+def make_summary_plot(it, raw, output, net_output, seeds, target):
+ """
+ This function create and save a summary figure
+ """
+ f, axarr = plt.subplots(2, 2, figsize=(8, 9.5))
+ f.suptitle("RW summary, Iteration: " + repr(it))
+
+ axarr[0, 0].set_title("Ground Truth Image")
+ axarr[0, 0].imshow(raw[0].detach().numpy(), cmap="gray")
+ axarr[0, 0].imshow(target[0, 0].detach().numpy(), alpha=0.6, vmin=-3, cmap="prism_r")
+ seeds_listx, seeds_listy = np.where(seeds[0].data != 0)
+ axarr[0, 0].scatter(seeds_listy,
+ seeds_listx, c="r")
+ axarr[0, 0].axis("off")
+
+ axarr[0, 1].set_title("LRW output (white seed)")
+ axarr[0, 1].imshow(raw[0].detach().numpy(), cmap="gray")
+ axarr[0, 1].imshow(np.argmax(output[0][0].detach().numpy(), 0), alpha=0.6, vmin=-3, cmap="prism_r")
+ axarr[0, 1].axis("off")
+
+ axarr[1, 0].set_title("Vertical Diffusivities")
+ axarr[1, 0].imshow(net_output[0, 0].detach().numpy(), cmap="gray")
+ axarr[1, 0].axis("off")
+
+ axarr[1, 1].set_title("Horizontal Diffusivities")
+ axarr[1, 1].imshow(net_output[0, 1].detach().numpy(), cmap="gray")
+ axarr[1, 1].axis("off")
+
+ plt.tight_layout()
+ plt.savefig("./results/%04i.png"%it)
+ plt.close()
+
+
+def draw_fragments_seeds(out_path, k, pred, pred_seed, gt, gt_seed, f_txt, raw=None, alpha=0.8):
+ m,n = pred.shape
+ ids = np.unique(pred)
+ size = len(ids)
+ print("k = %d, the neurons number of pred is %d" % (k, size))
+ f_txt.write("k = %d, the neurons number of pred is %d" % (k, size))
+ f_txt.write('\n')
+ color_pred = np.zeros([m, n, 3])
+ idx = np.searchsorted(ids, pred)
+ for i in range(3):
+ color_val = np.random.randint(0, 255, ids.shape)
+ if ids[0] == 0:
+ color_val[0] = 0
+ color_pred[:,:,i] = color_val[idx]
+ color_pred = color_pred / 255
+ if pred_seed is not None:
+ pred_seeds_listx, pred_seeds_listy = np.where(pred_seed != 0)
+
+ ids = np.unique(gt)
+ size = len(ids)
+ print("k = %d, the neurons number of gt is %d" % (k, size))
+ f_txt.write("k = %d, the neurons number of gt is %d" % (k, size))
+ f_txt.write('\n')
+ color_gt= np.zeros([m, n, 3])
+ idx = np.searchsorted(ids, gt)
+ for i in range(3):
+ color_val = np.random.randint(0, 255, ids.shape)
+ if ids[0] == 0:
+ color_val[0] = 0
+ color_gt[:,:,i] = color_val[idx]
+ color_gt = color_gt / 255
+ if gt_seed is not None:
+ gt_seeds_listx, gt_seeds_listy = np.where(gt_seed != 0)
+
+ if raw is not None:
+ plt.figure(figsize=(20,20),dpi=100)
+ plt.subplot(121)
+ plt.imshow(raw)
+ plt.imshow(color_pred, alpha=alpha)
+ if pred_seed is not None:
+ plt.scatter(pred_seeds_listy, pred_seeds_listx, c='k', marker='.')
+ plt.axis('off')
+ plt.subplot(122)
+ plt.imshow(raw)
+ plt.imshow(color_gt, alpha=alpha)
+ if gt_seed is not None:
+ plt.scatter(gt_seeds_listy, gt_seeds_listx, c='k', marker='.')
+ plt.axis('off')
+ # plt.show()
+ plt.savefig(os.path.join(out_path, str(k).zfill(4)+'.png'), bbox_inches = 'tight')
+ else:
+ plt.figure(figsize=(20,20),dpi=100)
+ plt.subplot(121)
+ plt.imshow(color_pred)
+ if pred_seed is not None:
+ plt.scatter(pred_seeds_listy, pred_seeds_listx, c='b')
+ plt.axis('off')
+ plt.subplot(122)
+ plt.imshow(color_gt)
+ if gt_seed is not None:
+ plt.scatter(gt_seeds_listy, gt_seeds_listx, c='b')
+ plt.axis('off')
+ # plt.show()
+ plt.savefig(os.path.join(out_path, str(k).zfill(4)+'.png'), bbox_inches = 'tight')
+ plt.close('all')
+
+
+def draw_fragments_noseeds(out_path, k, pred, gt=None, raw=None, alpha=0.8):
+ m,n = pred.shape
+ ids = np.unique(pred)
+ size = len(ids)
+ print("k = %d, the neurons number of pred is %d" % (k, size))
+ color_pred = np.zeros([m, n, 3])
+ idx = np.searchsorted(ids, pred)
+ for i in range(3):
+ color_val = np.random.randint(0, 255, ids.shape)
+ if ids[0] == 0:
+ color_val[0] = 0
+ color_pred[:,:,i] = color_val[idx]
+ color_pred = color_pred / 255
+
+ if gt is not None:
+ ids = np.unique(gt)
+ size = len(ids)
+ print("k = %d, the neurons number of gt is %d" % (k, size))
+ color_gt= np.zeros([m, n, 3])
+ idx = np.searchsorted(ids, gt)
+ for i in range(3):
+ color_val = np.random.randint(0, 255, ids.shape)
+ if ids[0] == 0:
+ color_val[0] = 0
+ color_gt[:,:,i] = color_val[idx]
+ color_gt = color_gt / 255
+
+ if gt is not None:
+ if raw is not None:
+ plt.figure(figsize=(20,20),dpi=100)
+ plt.subplot(121)
+ plt.imshow(raw)
+ plt.imshow(color_pred, alpha=alpha)
+ plt.axis('off')
+ plt.subplot(122)
+ plt.imshow(raw)
+ plt.imshow(color_gt, alpha=alpha)
+ plt.axis('off')
+ else:
+ plt.figure(figsize=(20,20),dpi=100)
+ plt.subplot(121)
+ plt.imshow(color_pred)
+ plt.axis('off')
+ plt.subplot(122)
+ plt.imshow(color_gt)
+ plt.axis('off')
+ else:
+ if raw is not None:
+ plt.figure(figsize=(20,20),dpi=100)
+ plt.imshow(raw)
+ plt.imshow(color_pred, alpha=alpha)
+ plt.axis('off')
+ else:
+ plt.figure(figsize=(20,20),dpi=100)
+ plt.imshow(color_pred)
+ plt.axis('off')
+ plt.savefig(os.path.join(out_path, str(k).zfill(4)+'.png'), bbox_inches = 'tight')
+ plt.close('all')
+
+
+def draw_fragments_3d(out_path, pred, gt=None, raw=None, alpha=0.8):
+ d,m,n = pred.shape
+ ids = np.unique(pred)
+ size = len(ids)
+ print("the neurons number of pred is %d" % size)
+ color_pred = np.zeros([d, m, n, 3])
+ idx = np.searchsorted(ids, pred)
+ for i in range(3):
+ color_val = np.random.randint(0, 255, ids.shape)
+ if ids[0] == 0:
+ color_val[0] = 0
+ color_pred[:,:,:,i] = color_val[idx]
+ color_pred = color_pred / 255
+
+ if gt is not None:
+ ids = np.unique(gt)
+ size = len(ids)
+ print("the neurons number of gt is %d" % size)
+ color_gt= np.zeros([d, m, n, 3])
+ idx = np.searchsorted(ids, gt)
+ for i in range(3):
+ color_val = np.random.randint(0, 255, ids.shape)
+ if ids[0] == 0:
+ color_val[0] = 0
+ color_gt[:,:,:,i] = color_val[idx]
+ color_gt = color_gt / 255
+
+ if gt is not None:
+ if raw is not None:
+ for k in range(d):
+ plt.figure(figsize=(20,20),dpi=100)
+ plt.subplot(121)
+ plt.imshow(raw[k])
+ plt.imshow(color_pred[k], alpha=alpha)
+ plt.axis('off')
+ plt.subplot(122)
+ plt.imshow(raw[k])
+ plt.imshow(color_gt[k], alpha=alpha)
+ plt.axis('off')
+ plt.savefig(os.path.join(out_path, str(k).zfill(4)+'.png'), bbox_inches = 'tight')
+ plt.close('all')
+ else:
+ for k in range(d):
+ plt.figure(figsize=(20,20),dpi=100)
+ plt.subplot(121)
+ plt.imshow(color_pred[k])
+ plt.axis('off')
+ plt.subplot(122)
+ plt.imshow(color_gt[k])
+ plt.axis('off')
+ plt.savefig(os.path.join(out_path, str(k).zfill(4)+'.png'), bbox_inches = 'tight')
+ plt.close('all')
+ else:
+ if raw is not None:
+ for k in range(d):
+ plt.figure(figsize=(20,20),dpi=100)
+ plt.imshow(raw[k])
+ plt.imshow(color_pred[k], alpha=alpha)
+ plt.axis('off')
+ plt.savefig(os.path.join(out_path, str(k).zfill(4)+'.png'), bbox_inches = 'tight')
+ plt.close('all')
+ else:
+ for k in range(d):
+ plt.figure(figsize=(20,20),dpi=100)
+ plt.imshow(color_pred[k])
+ plt.axis('off')
+ plt.savefig(os.path.join(out_path, str(k).zfill(4)+'.png'), bbox_inches = 'tight')
+ plt.close('all')
+
+
+if __name__ == "__main__":
+ in_path1 = '../data/snemi3d/AC4_inputs.h5'
+ in_path2 = '../data/snemi3d/AC4_labels.h5'
+ f = h5py.File(in_path1, 'r')
+ raw = f['main'][:]
+ f.close()
+
+ f = h5py.File(in_path2, 'r')
+ labels = f['main'][:]
+ f.close()
+
+ out_path = '../data/snemi3d/AC4'
+ if not os.path.exists(out_path):
+ os.mkdir(out_path)
+
+ draw_fragments_3d(out_path, labels, None, raw)
\ No newline at end of file
diff --git a/legacy/Train_and_Inference/utils/seg_util.py b/legacy/Train_and_Inference/utils/seg_util.py
new file mode 100644
index 0000000..3797c77
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/seg_util.py
@@ -0,0 +1,194 @@
+import numpy as np
+from scipy.sparse import coo_matrix
+from scipy.ndimage.morphology import binary_erosion,binary_dilation
+
+# reduce the labeling
+def relabel(segmentation):
+ # get the unique labels
+ uid = np.unique(segmentation)
+ # get the maximum label for the segment
+ mid = int(uid.max()) + 1
+
+ # create an array from original segment id to reduced id
+ mapping = np.zeros(mid, dtype=segmentation.dtype)
+ mapping[uid] = np.arange(len(uid), dtype=segmentation.dtype)
+ return mapping[segmentation]
+
+def remove_small(seg, thres=100):
+ sz = seg.shape
+ seg = seg.reshape(-1)
+ uid, uc = np.unique(seg, return_counts=True)
+ seg[np.in1d(seg,uid[uc0)
+ #stel=np.array([[1, 1],[1,1]]).astype(bool)
+ stel=np.array([[1,1,1], [1,1,1], [1,1,1]]).astype(bool)
+ #stel=np.array([[1,1,1,1],[1, 1, 1, 1],[1,1,1,1],[1,1,1,1]]).astype(bool)
+ gg3gd=np.zeros(gg3g.shape)
+ for i in range(gg3g.shape[0]):
+ gg3gd[i,:,:]=binary_dilation(gg3g[i,:,:],structure=stel,iterations=iter_num)
+ out = gg3.copy()
+ out[gg3gd==1]=0
+ return out
+
+def markInvalid(seg, iter_num=2, do_2d=True):
+ # find invalid
+ # if do erosion(seg==0), then miss the border
+ if do_2d:
+ stel=np.array([[1,1,1], [1,1,1]]).astype(bool)
+ if len(seg.shape)==2:
+ out = binary_dilation(seg>0, structure=stel, iterations=iter_num)
+ seg[out==0] = -1
+ else: # save memory
+ for z in range(seg.shape[0]):
+ tmp = seg[z] # by reference
+ out = binary_dilation(tmp>0, structure=stel, iterations=iter_num)
+ tmp[out==0] = -1
+ else:
+ stel=np.array([[1,1,1], [1,1,1], [1,1,1]]).astype(bool)
+ out = binary_dilation(seg>0, structure=stel, iterations=iter_num)
+ seg[out==0] = -1
+ return seg
diff --git a/legacy/Train_and_Inference/utils/show.py b/legacy/Train_and_Inference/utils/show.py
new file mode 100644
index 0000000..1f2103d
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/show.py
@@ -0,0 +1,273 @@
+import os
+import math
+import numpy as np
+from PIL import Image
+
+def show(img3d):
+ # only used for image with shape [18, 160, 160, 3]
+ num = img3d.shape[0]
+ column = 5
+ row = math.ceil(num / float(column))
+ size = img3d.shape[1]
+ img_all = np.zeros((size*row, size*column, 3), dtype=np.uint8)
+ for i in range(row):
+ for j in range(column):
+ index = i*column + j
+ if index >= num:
+ img = np.zeros_like(img3d[0], dtype=np.uint8)
+ else:
+ if np.max(img3d[index]) > 1:
+ img = (img3d[index]).astype(np.uint8)
+ else:
+ img = (img3d[index] * 255).astype(np.uint8)
+ img_all[i*size:(i+1)*size, j*size:(j+1)*size] = img
+ return img_all
+
+def show_one(img3d):
+ # only used for image with shape [18, 160, 160]
+ num = img3d.shape[0]
+ column = 5
+ row = math.ceil(num / float(column))
+ size = img3d.shape[1]
+ img_all = np.zeros((size*row, size*column), dtype=np.uint8)
+ for i in range(row):
+ for j in range(column):
+ index = i*column + j
+ if index >= num:
+ img = np.zeros_like(img3d[0], dtype=np.uint8)
+ else:
+ img = (img3d[index] * 255).astype(np.uint8)
+ img_all[i*size:(i+1)*size, j*size:(j+1)*size] = img
+ return img_all
+
+
+def show_one_(img3d):
+ num = img3d.shape[0]
+ column = 5
+ row = math.ceil(num / float(column))
+ size = img3d.shape[1]
+ img_all = np.zeros((size*row, size*column), dtype=np.uint8)
+ for i in range(row):
+ for j in range(column):
+ index = i*column + j
+ if index >= num:
+ img = np.zeros_like(img3d[0], dtype=np.uint8)
+ else:
+ img = (img3d[index]).astype(np.uint8)
+ img_all[i*size:(i+1)*size, j*size:(j+1)*size] = img
+ return img_all
+
+def show_CE(img3d):
+ # only used for image with shape [18, 160, 160]
+ num = img3d.shape[0]
+ column = 5
+ row = math.ceil(num / float(column))
+ size = img3d.shape[1]
+ img_all = np.zeros((size*row, size*column), dtype=np.uint8)
+ for i in range(row):
+ for j in range(column):
+ index = i*column + j
+ if index >= num:
+ img = np.zeros_like(img3d[0], dtype=np.uint8)
+ else:
+ img = (img3d[index]).astype(np.uint8)
+ img_all[i*size:(i+1)*size, j*size:(j+1)*size] = img
+ return img_all
+
+def training_show(iters, inputs, label, pred_bound, cache_path, if_skele=None, skele=None, pred_skele=None):
+ img_input = np.repeat(inputs[0].data.cpu().numpy(), 3, 0)
+ img_input = np.transpose(img_input, (1,2,3,0))
+ img_input = show(img_input)
+ input_placehplder = np.zeros_like(img_input, dtype=np.uint8)
+ im_cat1 = np.concatenate([img_input, input_placehplder], axis=1)
+
+ img_label = label[0][0:3].data.cpu().numpy()
+ img_label = np.transpose(img_label, (1,2,3,0))
+ img_label = show(img_label)
+
+ img_pred_bound = pred_bound[0][0:3].data.cpu().numpy()
+ img_pred_bound = np.transpose(img_pred_bound, (1,2,3,0))
+ img_pred_bound = show(img_pred_bound)
+ im_cat2 = np.concatenate([img_pred_bound, img_label], axis=1)
+
+ if if_skele is not None:
+ img_skele = np.repeat(skele[0, 0:1].data.cpu().numpy(), 3, 0)
+ img_skele = np.transpose(img_skele, (1,2,3,0))
+ img_skele = show(img_skele)
+
+ img_pred_skele = np.repeat(pred_skele[0, 0:1].data.cpu().numpy(), 3, 0)
+ img_pred_skele = np.transpose(img_pred_skele, (1,2,3,0))
+ img_pred_skele = show(img_pred_skele)
+ im_cat3 = np.concatenate([img_pred_skele, img_skele], axis=1)
+
+ im_cat = np.concatenate([im_cat1, im_cat2, im_cat3], axis=0)
+ else:
+ im_cat = np.concatenate([im_cat1, im_cat2], axis=0)
+ Image.fromarray(im_cat).save(os.path.join(cache_path, '%06d.png' % iters))
+
+def training_show_pretrain(iters, pred, label, cache_path, loss_mode='CrossEntropy'):
+ img_input = pred[0].data.cpu().numpy()
+ if loss_mode == 'CrossEntropy':
+ img_input = show_CE(img_input)
+ else:
+ img_input[img_input < 0] = 0
+ img_input[img_input > 1] = 1
+ img_input = show_one(img_input)
+ img_label = label[0].data.cpu().numpy()
+ img_label = show_one(img_label)
+ im_cat = np.concatenate([img_input, img_label], axis=1)
+ Image.fromarray(im_cat).save(os.path.join(cache_path, '%06d.png' % iters))
+
+
+def show_inpaining(iters, pred, label, mask, cache_path):
+ pred = pred[0].data.cpu().numpy()
+ label = label[0].data.cpu().numpy()
+ mask = mask[0].data.cpu().numpy()
+ inputs = label * mask
+ inputs = np.squeeze(inputs)
+ pred = np.squeeze(pred)
+ inputs = inputs[14:-14, 106:-106, 106:-106]
+ pred[pred < 0] = 0; pred[pred > 1] =1
+ pred_img = show_one(pred)
+ inputs_img = show_one(inputs)
+ im_cat = np.concatenate([inputs_img, pred_img], axis=1)
+ Image.fromarray(im_cat).save(os.path.join(cache_path, '%06d.png' % iters))
+
+
+def show_affs(iters, inputs, pred, target, cache_path, model_type='mala'):
+ pred = pred[0].data.cpu().numpy()
+ inputs = inputs[0].data.cpu().numpy()
+ target = target[0].data.cpu().numpy()
+ inputs = np.squeeze(inputs)
+ if model_type == 'mala':
+ inputs = inputs[14:-14, 106:-106, 106:-106]
+ inputs = inputs[:,:,:,np.newaxis]
+ inputs = np.repeat(inputs, 3, 3)
+ pred = np.transpose(pred, (1,2,3,0))
+ target = np.transpose(target, (1,2,3,0))
+ inputs[inputs<0]=0; inputs[inputs>1]=1
+ pred[pred<0]=0; pred[pred>1]=1
+ target[target<0]=0; target[target>1]=1
+ inputs_img = show(inputs)
+ pred_img = show(pred)
+ target_img = show(target)
+ im_cat = np.concatenate([inputs_img, pred_img, target_img], axis=1)
+ Image.fromarray(im_cat).save(os.path.join(cache_path, '%06d.png' % iters))
+
+def show_bound(iters, inputs, pred, target, cache_path, model_type='mala'):
+ pred = pred[0].data.cpu().numpy()
+ inputs = inputs[0].data.cpu().numpy()
+ target = target[0].data.cpu().numpy()
+
+ inputs = np.squeeze(inputs)
+ inputs = inputs[:,:,:,np.newaxis]
+ inputs = np.repeat(inputs, 3, 3)
+
+ pred = np.squeeze(pred)
+ pred = pred[:,:,:,np.newaxis]
+ pred = np.repeat(pred, 3, 3)
+
+ target = np.squeeze(target)
+ target = target[:,:,:,np.newaxis]
+ target = np.repeat(target, 3, 3)
+
+ inputs[inputs<0]=0; inputs[inputs>1]=1
+ pred[pred<0]=0; pred[pred>1]=1
+ target[target<0]=0; target[target>1]=1
+ inputs_img = show(inputs)
+ pred_img = show(pred)
+ target_img = show(target)
+ im_cat = np.concatenate([inputs_img, pred_img, target_img], axis=1)
+ Image.fromarray(im_cat).save(os.path.join(cache_path, '%06d.png' % iters))
+
+def class_color(lb):
+ d, h, w = lb.shape
+ lb_color1 = np.zeros((d, h, w), dtype=np.uint8)
+ lb_color2 = np.zeros((d, h, w), dtype=np.uint8)
+ lb_color3 = np.zeros((d, h, w), dtype=np.uint8)
+ ids_0 = lb == 0
+ ids_1 = lb == 1
+ lb_color1[ids_0] = 0; lb_color2[ids_0] = 0; lb_color3[ids_0] = 255
+ lb_color1[ids_1] = 0; lb_color2[ids_1] = 255; lb_color3[ids_1] = 0
+ lb_color = np.concatenate([lb_color1[:,:,:,np.newaxis], lb_color2[:,:,:,np.newaxis], lb_color3[:,:,:,np.newaxis]], axis=3)
+ return lb_color
+
+
+def show_affs_pseudo(iters, inputs, pred, target, mask, cache_path, model_type='mala'):
+ pred = pred[0].data.cpu().numpy()
+ inputs = inputs[0].data.cpu().numpy()
+ target = target[0].data.cpu().numpy()
+ mask = mask[0].data.cpu().numpy()
+ inputs = np.squeeze(inputs)
+ if model_type == 'mala':
+ inputs = inputs[14:-14, 106:-106, 106:-106]
+ inputs = inputs[:,:,:,np.newaxis]
+ inputs = np.repeat(inputs, 3, 3)
+ pred = np.transpose(pred, (1,2,3,0))
+ target = np.transpose(target, (1,2,3,0))
+ affs_z = class_color(target[:, :, :, 0]) * mask[0][:,:,:,np.newaxis]
+ affs_y = class_color(target[:, :, :, 1]) * mask[1][:,:,:,np.newaxis]
+ affs_x = class_color(target[:, :, :, 2]) * mask[2][:,:,:,np.newaxis]
+ inputs_img = show(inputs)
+ pred_img = show(pred)
+ # target_img = show(target)
+ mask = np.transpose(mask, (1,2,3,0))
+ mask_img = show(mask)
+ affs_z_img = show(affs_z)
+ affs_y_img = show(affs_y)
+ affs_x_img = show(affs_x)
+ # im_cat = np.concatenate([inputs_img, pred_img, target_img], axis=1)
+ im_cat1 = np.concatenate([inputs_img, pred_img], axis=1)
+ im_cat2 = np.concatenate([mask_img, affs_z_img], axis=1)
+ im_cat3 = np.concatenate([affs_y_img, affs_x_img], axis=1)
+ im_cat = np.concatenate([im_cat1, im_cat2, im_cat3], axis=0)
+ Image.fromarray(im_cat).save(os.path.join(cache_path, '%06d.png' % iters))
+
+def show_affs_whole(iters, out_affs, gt_affs, cache_path, index):
+ out_affs = out_affs[:, -1, ...]
+ gt_affs = gt_affs[:, -1, ...]
+ out_affs = (out_affs * 255).astype(np.uint8)
+ out_affs = np.transpose(out_affs, (1,2,0))
+ gt_affs = (gt_affs * 255).astype(np.uint8)
+ gt_affs = np.transpose(gt_affs, (1,2,0))
+ im_cat = np.concatenate([out_affs, gt_affs], axis=1)
+ Image.fromarray(im_cat).save(os.path.join(cache_path, '%06d_%d.png' % (iters,index)))
+
+def show_bound_whole(iters, out_affs, gt_affs, cache_path, index):
+ out_affs = out_affs.squeeze()[0]
+ out_affs = (out_affs * 255).astype(np.uint8)
+ gt_affs = gt_affs.squeeze()[0]
+ gt_affs = (gt_affs * 255).astype(np.uint8)
+ im_cat = np.concatenate([out_affs, gt_affs], axis=1)
+ Image.fromarray(im_cat).save(os.path.join(cache_path, '%06d_%d.png' % (iters,index)))
+
+def show_affs_consistency(iters, inputs, pred, target, inputs_u, out_u1, out_u2, cache_path):
+ pred = pred[0].data.cpu().numpy()
+ inputs = inputs[0].data.cpu().numpy()
+ target = target[0].data.cpu().numpy()
+ inputs = np.squeeze(inputs)
+ inputs = inputs[14:-14, 106:-106, 106:-106]
+ inputs = inputs[:,:,:,np.newaxis]
+ inputs = np.repeat(inputs, 3, 3)
+ pred = np.transpose(pred, (1,2,3,0))
+ target = np.transpose(target, (1,2,3,0))
+ inputs_img = show(inputs)
+ pred_img = show(pred)
+ target_img = show(target)
+ im_cat1 = np.concatenate([inputs_img, pred_img, target_img], axis=1)
+
+ out_u1 = out_u1[0].data.cpu().numpy()
+ inputs_u = inputs_u[0].data.cpu().numpy()
+ out_u2 = out_u2[0].data.cpu().numpy()
+ inputs_u = np.squeeze(inputs_u)
+ inputs_u = inputs_u[14:-14, 106:-106, 106:-106]
+ inputs_u = inputs_u[:,:,:,np.newaxis]
+ inputs_u = np.repeat(inputs_u, 3, 3)
+ out_u1 = np.transpose(out_u1, (1,2,3,0))
+ out_u2 = np.transpose(out_u2, (1,2,3,0))
+ inputs_u_img = show(inputs_u)
+ out_u1_img = show(out_u1)
+ out_u2_img = show(out_u2)
+ im_cat2 = np.concatenate([inputs_u_img, out_u1_img, out_u2_img], axis=1)
+ im_cat = np.concatenate([im_cat1, im_cat2], axis=0)
+ Image.fromarray(im_cat).save(os.path.join(cache_path, '%06d.png' % iters))
\ No newline at end of file
diff --git a/legacy/Train_and_Inference/utils/torch_utils.py b/legacy/Train_and_Inference/utils/torch_utils.py
new file mode 100644
index 0000000..b56b616
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/torch_utils.py
@@ -0,0 +1,17 @@
+from distutils.version import LooseVersion, StrictVersion
+import torch
+
+# align_corners option available for v1.3.0 and above
+HAS_AFFINE_ALIGN_CORNERS = LooseVersion(torch.__version__) >= LooseVersion('1.3.0')
+# align_corners defaults to True before v1.4.0, False from v1.4.0 and after
+AFFINE_ALIGN_CORNERS_DEFAULT = LooseVersion(torch.__version__) <= LooseVersion('1.3.0')
+
+
+def affine_align_corners_kw(val):
+ if HAS_AFFINE_ALIGN_CORNERS:
+ return dict(align_corners=val)
+ else:
+ if not val:
+ raise RuntimeError('align_corners not available in torch version {} so '
+ 'cannot set to False'.format(torch.__version__))
+ return {}
diff --git a/legacy/Train_and_Inference/utils/utils.py b/legacy/Train_and_Inference/utils/utils.py
new file mode 100644
index 0000000..6fc3be1
--- /dev/null
+++ b/legacy/Train_and_Inference/utils/utils.py
@@ -0,0 +1,57 @@
+import torch
+import random
+import numpy as np
+import subprocess
+
+def setup_seed(seed):
+ torch.manual_seed(seed)
+ # torch.cuda.manual_seed(seed)
+ torch.cuda.manual_seed_all(seed)
+ np.random.seed(seed)
+ random.seed(seed)
+ torch.backends.cudnn.deterministic = True
+
+def execute(cmd):
+ popen = subprocess.Popen(cmd, stdout=subprocess.PIPE, universal_newlines=True)
+ for stdout_line in iter(popen.stdout.readline, ""):
+ yield stdout_line
+ popen.stdout.close()
+ return_code = popen.wait()
+ if return_code:
+ raise subprocess.CalledProcessError(return_code, cmd)
+
+def compute_num_single(size, stride):
+ # 计算大概需要滑窗的个数
+ num_window = size // stride
+ # 判断是否能整除，如果不可以就+1
+ # 即滑窗个数加1
+ if size % stride != 0:
+ num_window += 1
+ # 计算padding的大小
+ padding_2times = num_window * stride - size + (stride * 2)
+ # 如果padding的大小不能被2整除
+ # 就继续增加滑窗的个数，已使得padding_2times能被2整除
+ while padding_2times % 2 != 0:
+ num_window += 1
+ padding_2times = num_window * stride - size + (stride * 2)
+ # 除以2是为了对称padding
+ padding = padding_2times // 2
+ # 加1是为了补上最后一个完整的滑窗
+ num_window += 1
+ return num_window, padding
+
+def compute_num(raw_shape, stride):
+ size_z = raw_shape[0]
+ size_xy = raw_shape[1]
+ stride_z = stride[0]
+ stride_xy = stride[1]
+ num_z, padding_z = compute_num_single(size_z, stride_z)
+ num_xy, padding_xy = compute_num_single(size_xy, stride_xy)
+ return [num_z, num_xy, num_xy], [padding_z, padding_xy, padding_xy]
+
+if __name__ == "__main__":
+ raw = [500, 4096, 4096]
+ stride = [18, 128, 128]
+ num, padding = compute_num(raw, stride)
+ print(num)
+ print(padding)
diff --git a/legacy/requirements.txt b/legacy/requirements.txt
new file mode 100644
index 0000000..98c7bf8
--- /dev/null
+++ b/legacy/requirements.txt
@@ -0,0 +1,248 @@
+addict==2.4.0
+asttokens==2.2.1
+atomicwrites==1.4.1
+attrdict==2.0.1
+attrs==23.1.0
+backcall==0.2.0
+blinker==1.6.2
+boto3==1.26.118
+botocore==1.29.118
+Brotli==1.0.9
+cached-property @ file:///home/conda/feedstock_root/build_artifacts/cached_property_1615209429212/work
+cachetools==5.3.0
+certifi==2023.5.7
+cffi==1.15.1
+chardet==5.1.0
+charset-normalizer @ file:///home/conda/feedstock_root/build_artifacts/charset-normalizer_1678108872112/work
+click @ file:///home/conda/feedstock_root/build_artifacts/click_1666798198223/work
+cloud-files==4.15.1
+cloud-volume==8.19.3
+cloudpickle @ file:///home/conda/feedstock_root/build_artifacts/cloudpickle_1674202310934/work
+cmake==3.26.3
+colorama @ file:///home/conda/feedstock_root/build_artifacts/colorama_1666700638685/work
+coloredlogs==15.0.1
+colorlog==6.7.0
+comm==0.1.3
+compressed-segmentation==2.2.0
+compresso==3.2.0
+ConfigArgParse==1.5.3
+connected-components-3d==3.10.5
+contourpy @ file:///home/conda/feedstock_root/build_artifacts/contourpy_1686733808435/work
+crackle-codec==0.6.1
+crc32c==2.3.post0
+cryptography==40.0.2
+cycler @ file:///home/conda/feedstock_root/build_artifacts/cycler_1635519461629/work
+Cython==0.29.34
+cytoolz @ file:///home/conda/feedstock_root/build_artifacts/cytoolz_1666829688252/work
+dash==2.9.3
+dash-core-components==2.0.0
+dash-html-components==2.0.0
+dash-table==5.0.0
+dask @ file:///home/conda/feedstock_root/build_artifacts/dask-core_1687813538832/work
+debugpy==1.6.7
+decorator==5.1.1
+deflate==0.3.0
+dijkstra3d==1.12.0
+dill==0.3.6
+DracoPy==1.0.1
+edt==2.3.0
+einops==0.8.0
+elf @ file:///home/conda/feedstock_root/build_artifacts/python-elf_1684541229700/work
+et-xmlfile @ file:///home/conda/feedstock_root/build_artifacts/et_xmlfile_1674664118162/work
+exceptiongroup==1.1.1
+executing==1.2.0
+fasteners==0.18
+fastjsonschema==2.16.3
+fastremap==1.13.4
+filelock==3.12.0
+fill-voids==2.0.3
+Flask==2.3.2
+flatbuffers==23.3.3
+fonttools @ file:///home/conda/feedstock_root/build_artifacts/fonttools_1686578436974/work
+fpzip==1.2.0
+fsspec @ file:///home/conda/feedstock_root/build_artifacts/fsspec_1686342280219/work
+gevent==22.10.2
+google-api-core==2.11.0
+google-apitools==0.5.32
+google-auth==2.17.3
+google-cloud-appengine-logging==1.3.0
+google-cloud-audit-log==0.2.5
+google-cloud-core==2.3.2
+google-cloud-logging==3.5.0
+google-cloud-storage==2.8.0
+google-crc32c==1.5.0
+google-resumable-media==2.4.1
+googleapis-common-protos==1.59.0
+greenlet==2.0.2
+grpc-google-iam-v1==0.12.6
+grpcio==1.54.0
+grpcio-status==1.54.0
+h5py @ file:///home/conda/feedstock_root/build_artifacts/h5py_1687263083757/work
+httplib2==0.22.0
+huggingface-hub==0.23.4
+humanfriendly==10.0
+idna @ file:///home/conda/feedstock_root/build_artifacts/idna_1663625384323/work
+igneous-pipeline==3.0.0
+imagecodecs @ file:///home/conda/feedstock_root/build_artifacts/imagecodecs_1679699549914/work
+imageio @ file:///home/conda/feedstock_root/build_artifacts/imageio_1686552404404/work
+importlib-metadata @ file:///home/conda/feedstock_root/build_artifacts/importlib-metadata_1687138353019/work
+importlib-resources @ file:///home/conda/feedstock_root/build_artifacts/importlib_resources_1676919000169/work
+inflection==0.5.1
+iniconfig==2.0.0
+ipykernel==6.22.0
+ipython==8.13.1
+ipywidgets==8.0.6
+itsdangerous==2.1.2
+jedi==0.18.2
+Jinja2==3.1.2
+jmespath==1.0.1
+joblib @ file:///home/conda/feedstock_root/build_artifacts/joblib_1663332044897/work
+json5==0.9.11
+jsonschema==4.17.3
+jupyter_client==8.2.0
+jupyter_core==5.3.0
+jupyterlab-widgets==3.0.7
+kimimaro==3.3.0
+kiwisolver @ file:///home/conda/feedstock_root/build_artifacts/kiwisolver_1666805784128/work
+lazy_loader==0.2
+lit==16.0.1
+llvmlite==0.40.1
+locket @ file:///home/conda/feedstock_root/build_artifacts/locket_1650660393415/work
+lxml==4.9.2
+mapbox-earcut==1.0.1
+mapbuffer==0.7.0
+Markdown==3.4.3
+MarkupSafe==2.1.2
+matplotlib @ file:///home/conda/feedstock_root/build_artifacts/matplotlib-suite_1678135567769/work
+matplotlib-inline==0.1.6
+monai==1.1.0
+mpmath==1.3.0
+mrcfile @ file:///home/conda/feedstock_root/build_artifacts/mrcfile_1663870629802/work
+multiprocess==0.70.14
+munkres==1.1.4
+nbformat==5.7.0
+nest-asyncio==1.5.6
+networkx @ file:///home/conda/feedstock_root/build_artifacts/networkx_1680692919326/work
+neuroglancer==2.36
+nibabel==5.2.1
+numba @ file:///home/conda/feedstock_root/build_artifacts/numba_1687804907391/work
+numpy @ file:///home/conda/feedstock_root/build_artifacts/numpy_1687808300586/work
+nvidia-cublas-cu11==11.10.3.66
+nvidia-cuda-cupti-cu11==11.7.101
+nvidia-cuda-nvrtc-cu11==11.7.99
+nvidia-cuda-runtime-cu11==11.7.99
+nvidia-cudnn-cu11==8.5.0.96
+nvidia-cufft-cu11==10.9.0.58
+nvidia-curand-cu11==10.2.10.91
+nvidia-cusolver-cu11==11.4.0.1
+nvidia-cusparse-cu11==11.7.4.91
+nvidia-nccl-cu11==2.14.3
+nvidia-nvtx-cu11==11.7.91
+oauth2client==4.1.3
+onnx==1.13.1
+onnxruntime==1.14.1
+open3d==0.17.0
+opencv-python==4.7.0.72
+openmesh==1.2.1
+openpyxl @ file:///home/conda/feedstock_root/build_artifacts/openpyxl_1682610773138/work
+orjson==3.8.10
+packaging @ file:///home/conda/feedstock_root/build_artifacts/packaging_1681337016113/work
+pandas @ file:///home/conda/feedstock_root/build_artifacts/pandas_1685342911872/work
+parso==0.8.3
+partd @ file:///home/conda/feedstock_root/build_artifacts/partd_1681246756246/work
+pathos==0.3.0
+pbr==5.11.1
+pexpect==4.8.0
+pickleshare==0.7.5
+Pillow @ file:///home/conda/feedstock_root/build_artifacts/pillow_1684654072636/work
+platformdirs @ file:///home/conda/feedstock_root/build_artifacts/platformdirs_1687705014305/work
+plotly==5.14.1
+pluggy==1.0.0
+pooch @ file:///home/conda/feedstock_root/build_artifacts/pooch_1679580333621/work
+posix-ipc==1.1.1
+pox==0.3.2
+ppft==1.7.6.6
+prompt-toolkit==3.0.38
+proto-plus==1.22.2
+protobuf==3.20.1
+psutil==5.9.5
+ptyprocess==0.7.0
+pure-eval==0.2.2
+pyasn1==0.5.0
+pyasn1-modules==0.3.0
+pybind11==2.10.4
+pycocotools==2.0.6
+pycollada==0.7.2
+pycparser==2.21
+pyfqmr==0.1.2
+Pygments==2.15.1
+pynndescent==0.5.10
+pyOpenSSL==23.1.1
+pyparsing @ file:///home/conda/feedstock_root/build_artifacts/pyparsing_1687132014935/work
+pyquaternion==0.9.9
+pyrsistent==0.19.3
+pysimdjson==5.0.2
+PySocks @ file:///home/conda/feedstock_root/build_artifacts/pysocks_1661604839144/work
+pyspng-seunglab==1.0.0
+pytest==7.3.1
+python-dateutil @ file:///home/conda/feedstock_root/build_artifacts/python-dateutil_1626286286081/work
+python-jsonschema-objects==0.4.1
+pytz @ file:///home/conda/feedstock_root/build_artifacts/pytz_1680088766131/work
+PyWavelets @ file:///home/conda/feedstock_root/build_artifacts/pywavelets_1673082327051/work
+PyYAML @ file:///home/conda/feedstock_root/build_artifacts/pyyaml_1666772387118/work
+pyzmq==25.0.2
+quad-mesh-simplify==1.1.5
+requests @ file:///home/conda/feedstock_root/build_artifacts/requests_1684774241324/work
+rsa==4.9
+Rtree==1.0.1
+s3transfer==0.6.0
+safetensors==0.4.3
+scikit-image==0.20.0
+scikit-learn @ file:///home/conda/feedstock_root/build_artifacts/scikit-learn_1685023709438/work
+scipy @ file:///home/conda/feedstock_root/build_artifacts/scipy-split_1687763396626/work/base/dist/scipy-1.11.0-cp39-cp39-linux_x86_64.whl#sha256=5a44914d9d6aee4a88276a85f8fc857fdf47b30ccbaef31873614203f63769cf
+seaborn==0.12.2
+shapely==2.0.1
+shard-computer==1.1.0
+simplejpeg==1.6.3
+six @ file:///home/conda/feedstock_root/build_artifacts/six_1620240208055/work
+skan @ file:///home/conda/feedstock_root/build_artifacts/skan_1687524172302/work
+stack-data==0.6.2
+svg.path==6.2
+sympy==1.11.1
+task-queue==2.12.1
+tenacity==8.2.2
+tensorboardX==2.6
+thop==0.1.1.post2209072238
+threadpoolctl @ file:///home/conda/feedstock_root/build_artifacts/threadpoolctl_1643647933166/work
+tifffile @ file:///home/conda/feedstock_root/build_artifacts/tifffile_1681364231336/work
+timm==1.0.7
+tinybrain==1.3.1
+tomli==2.0.1
+toolz @ file:///home/conda/feedstock_root/build_artifacts/toolz_1657485559105/work
+torch==2.0.0
+torchaudio==2.0.1
+torchsummary==1.5.1
+torchvision==0.15.1
+tornado==6.3.1
+tqdm @ file:///home/conda/feedstock_root/build_artifacts/tqdm_1677948868469/work
+traitlets==5.9.0
+trimesh==3.21.5
+triton==2.0.0
+typing_extensions @ file:///home/conda/feedstock_root/build_artifacts/typing_extensions_1685704949284/work
+tzdata @ file:///home/conda/feedstock_root/build_artifacts/python-tzdata_1680081134351/work
+umap-learn==0.5.3
+unicodedata2 @ file:///home/conda/feedstock_root/build_artifacts/unicodedata2_1667239485250/work
+urllib3==1.26.16
+urllib3-secure-extra==0.1.0
+vit-pytorch==1.7.0
+wcwidth==0.2.6
+Werkzeug==2.3.3
+widgetsnbextension==4.0.7
+xxhash==3.2.0
+zfpc==0.1.2
+zfpy==1.0.0
+zipp @ file:///home/conda/feedstock_root/build_artifacts/zipp_1677313463193/work
+zmesh==1.0.0
+zope.event==4.6
+zope.interface==6.0
+zstandard==0.21.0
diff --git a/requirements-training.txt b/requirements-training.txt
new file mode 100644
index 0000000..9c884b9
--- /dev/null
+++ b/requirements-training.txt
@@ -0,0 +1,13 @@
+# Additional dependencies for the original CUDA research training entry points.
+-r requirements.txt
+numpy>=1.24,<2
+addict==2.4.0
+PyYAML>=6,<7
+tensorboardX>=2.6,<3
+einops>=0.8,<1
+imageio>=2.31,<3
+Pillow>=9.5
+scipy>=1.10,<2
+scikit-image>=0.20,<1
+opencv-python-headless>=4.7,<4.12
+matplotlib>=3.7,<4
diff --git a/requirements.txt b/requirements.txt
index 98c7bf8..e2a441d 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,248 +1,5 @@
-addict==2.4.0
-asttokens==2.2.1
-atomicwrites==1.4.1
-attrdict==2.0.1
-attrs==23.1.0
-backcall==0.2.0
-blinker==1.6.2
-boto3==1.26.118
-botocore==1.29.118
-Brotli==1.0.9
-cached-property @ file:///home/conda/feedstock_root/build_artifacts/cached_property_1615209429212/work
-cachetools==5.3.0
-certifi==2023.5.7
-cffi==1.15.1
-chardet==5.1.0
-charset-normalizer @ file:///home/conda/feedstock_root/build_artifacts/charset-normalizer_1678108872112/work
-click @ file:///home/conda/feedstock_root/build_artifacts/click_1666798198223/work
-cloud-files==4.15.1
-cloud-volume==8.19.3
-cloudpickle @ file:///home/conda/feedstock_root/build_artifacts/cloudpickle_1674202310934/work
-cmake==3.26.3
-colorama @ file:///home/conda/feedstock_root/build_artifacts/colorama_1666700638685/work
-coloredlogs==15.0.1
-colorlog==6.7.0
-comm==0.1.3
-compressed-segmentation==2.2.0
-compresso==3.2.0
-ConfigArgParse==1.5.3
-connected-components-3d==3.10.5
-contourpy @ file:///home/conda/feedstock_root/build_artifacts/contourpy_1686733808435/work
-crackle-codec==0.6.1
-crc32c==2.3.post0
-cryptography==40.0.2
-cycler @ file:///home/conda/feedstock_root/build_artifacts/cycler_1635519461629/work
-Cython==0.29.34
-cytoolz @ file:///home/conda/feedstock_root/build_artifacts/cytoolz_1666829688252/work
-dash==2.9.3
-dash-core-components==2.0.0
-dash-html-components==2.0.0
-dash-table==5.0.0
-dask @ file:///home/conda/feedstock_root/build_artifacts/dask-core_1687813538832/work
-debugpy==1.6.7
-decorator==5.1.1
-deflate==0.3.0
-dijkstra3d==1.12.0
-dill==0.3.6
-DracoPy==1.0.1
-edt==2.3.0
-einops==0.8.0
-elf @ file:///home/conda/feedstock_root/build_artifacts/python-elf_1684541229700/work
-et-xmlfile @ file:///home/conda/feedstock_root/build_artifacts/et_xmlfile_1674664118162/work
-exceptiongroup==1.1.1
-executing==1.2.0
-fasteners==0.18
-fastjsonschema==2.16.3
-fastremap==1.13.4
-filelock==3.12.0
-fill-voids==2.0.3
-Flask==2.3.2
-flatbuffers==23.3.3
-fonttools @ file:///home/conda/feedstock_root/build_artifacts/fonttools_1686578436974/work
-fpzip==1.2.0
-fsspec @ file:///home/conda/feedstock_root/build_artifacts/fsspec_1686342280219/work
-gevent==22.10.2
-google-api-core==2.11.0
-google-apitools==0.5.32
-google-auth==2.17.3
-google-cloud-appengine-logging==1.3.0
-google-cloud-audit-log==0.2.5
-google-cloud-core==2.3.2
-google-cloud-logging==3.5.0
-google-cloud-storage==2.8.0
-google-crc32c==1.5.0
-google-resumable-media==2.4.1
-googleapis-common-protos==1.59.0
-greenlet==2.0.2
-grpc-google-iam-v1==0.12.6
-grpcio==1.54.0
-grpcio-status==1.54.0
-h5py @ file:///home/conda/feedstock_root/build_artifacts/h5py_1687263083757/work
-httplib2==0.22.0
-huggingface-hub==0.23.4
-humanfriendly==10.0
-idna @ file:///home/conda/feedstock_root/build_artifacts/idna_1663625384323/work
-igneous-pipeline==3.0.0
-imagecodecs @ file:///home/conda/feedstock_root/build_artifacts/imagecodecs_1679699549914/work
-imageio @ file:///home/conda/feedstock_root/build_artifacts/imageio_1686552404404/work
-importlib-metadata @ file:///home/conda/feedstock_root/build_artifacts/importlib-metadata_1687138353019/work
-importlib-resources @ file:///home/conda/feedstock_root/build_artifacts/importlib_resources_1676919000169/work
-inflection==0.5.1
-iniconfig==2.0.0
-ipykernel==6.22.0
-ipython==8.13.1
-ipywidgets==8.0.6
-itsdangerous==2.1.2
-jedi==0.18.2
-Jinja2==3.1.2
-jmespath==1.0.1
-joblib @ file:///home/conda/feedstock_root/build_artifacts/joblib_1663332044897/work
-json5==0.9.11
-jsonschema==4.17.3
-jupyter_client==8.2.0
-jupyter_core==5.3.0
-jupyterlab-widgets==3.0.7
-kimimaro==3.3.0
-kiwisolver @ file:///home/conda/feedstock_root/build_artifacts/kiwisolver_1666805784128/work
-lazy_loader==0.2
-lit==16.0.1
-llvmlite==0.40.1
-locket @ file:///home/conda/feedstock_root/build_artifacts/locket_1650660393415/work
-lxml==4.9.2
-mapbox-earcut==1.0.1
-mapbuffer==0.7.0
-Markdown==3.4.3
-MarkupSafe==2.1.2
-matplotlib @ file:///home/conda/feedstock_root/build_artifacts/matplotlib-suite_1678135567769/work
-matplotlib-inline==0.1.6
-monai==1.1.0
-mpmath==1.3.0
-mrcfile @ file:///home/conda/feedstock_root/build_artifacts/mrcfile_1663870629802/work
-multiprocess==0.70.14
-munkres==1.1.4
-nbformat==5.7.0
-nest-asyncio==1.5.6
-networkx @ file:///home/conda/feedstock_root/build_artifacts/networkx_1680692919326/work
-neuroglancer==2.36
-nibabel==5.2.1
-numba @ file:///home/conda/feedstock_root/build_artifacts/numba_1687804907391/work
-numpy @ file:///home/conda/feedstock_root/build_artifacts/numpy_1687808300586/work
-nvidia-cublas-cu11==11.10.3.66
-nvidia-cuda-cupti-cu11==11.7.101
-nvidia-cuda-nvrtc-cu11==11.7.99
-nvidia-cuda-runtime-cu11==11.7.99
-nvidia-cudnn-cu11==8.5.0.96
-nvidia-cufft-cu11==10.9.0.58
-nvidia-curand-cu11==10.2.10.91
-nvidia-cusolver-cu11==11.4.0.1
-nvidia-cusparse-cu11==11.7.4.91
-nvidia-nccl-cu11==2.14.3
-nvidia-nvtx-cu11==11.7.91
-oauth2client==4.1.3
-onnx==1.13.1
-onnxruntime==1.14.1
-open3d==0.17.0
-opencv-python==4.7.0.72
-openmesh==1.2.1
-openpyxl @ file:///home/conda/feedstock_root/build_artifacts/openpyxl_1682610773138/work
-orjson==3.8.10
-packaging @ file:///home/conda/feedstock_root/build_artifacts/packaging_1681337016113/work
-pandas @ file:///home/conda/feedstock_root/build_artifacts/pandas_1685342911872/work
-parso==0.8.3
-partd @ file:///home/conda/feedstock_root/build_artifacts/partd_1681246756246/work
-pathos==0.3.0
-pbr==5.11.1
-pexpect==4.8.0
-pickleshare==0.7.5
-Pillow @ file:///home/conda/feedstock_root/build_artifacts/pillow_1684654072636/work
-platformdirs @ file:///home/conda/feedstock_root/build_artifacts/platformdirs_1687705014305/work
-plotly==5.14.1
-pluggy==1.0.0
-pooch @ file:///home/conda/feedstock_root/build_artifacts/pooch_1679580333621/work
-posix-ipc==1.1.1
-pox==0.3.2
-ppft==1.7.6.6
-prompt-toolkit==3.0.38
-proto-plus==1.22.2
-protobuf==3.20.1
-psutil==5.9.5
-ptyprocess==0.7.0
-pure-eval==0.2.2
-pyasn1==0.5.0
-pyasn1-modules==0.3.0
-pybind11==2.10.4
-pycocotools==2.0.6
-pycollada==0.7.2
-pycparser==2.21
-pyfqmr==0.1.2
-Pygments==2.15.1
-pynndescent==0.5.10
-pyOpenSSL==23.1.1
-pyparsing @ file:///home/conda/feedstock_root/build_artifacts/pyparsing_1687132014935/work
-pyquaternion==0.9.9
-pyrsistent==0.19.3
-pysimdjson==5.0.2
-PySocks @ file:///home/conda/feedstock_root/build_artifacts/pysocks_1661604839144/work
-pyspng-seunglab==1.0.0
-pytest==7.3.1
-python-dateutil @ file:///home/conda/feedstock_root/build_artifacts/python-dateutil_1626286286081/work
-python-jsonschema-objects==0.4.1
-pytz @ file:///home/conda/feedstock_root/build_artifacts/pytz_1680088766131/work
-PyWavelets @ file:///home/conda/feedstock_root/build_artifacts/pywavelets_1673082327051/work
-PyYAML @ file:///home/conda/feedstock_root/build_artifacts/pyyaml_1666772387118/work
-pyzmq==25.0.2
-quad-mesh-simplify==1.1.5
-requests @ file:///home/conda/feedstock_root/build_artifacts/requests_1684774241324/work
-rsa==4.9
-Rtree==1.0.1
-s3transfer==0.6.0
-safetensors==0.4.3
-scikit-image==0.20.0
-scikit-learn @ file:///home/conda/feedstock_root/build_artifacts/scikit-learn_1685023709438/work
-scipy @ file:///home/conda/feedstock_root/build_artifacts/scipy-split_1687763396626/work/base/dist/scipy-1.11.0-cp39-cp39-linux_x86_64.whl#sha256=5a44914d9d6aee4a88276a85f8fc857fdf47b30ccbaef31873614203f63769cf
-seaborn==0.12.2
-shapely==2.0.1
-shard-computer==1.1.0
-simplejpeg==1.6.3
-six @ file:///home/conda/feedstock_root/build_artifacts/six_1620240208055/work
-skan @ file:///home/conda/feedstock_root/build_artifacts/skan_1687524172302/work
-stack-data==0.6.2
-svg.path==6.2
-sympy==1.11.1
-task-queue==2.12.1
-tenacity==8.2.2
-tensorboardX==2.6
-thop==0.1.1.post2209072238
-threadpoolctl @ file:///home/conda/feedstock_root/build_artifacts/threadpoolctl_1643647933166/work
-tifffile @ file:///home/conda/feedstock_root/build_artifacts/tifffile_1681364231336/work
-timm==1.0.7
-tinybrain==1.3.1
-tomli==2.0.1
-toolz @ file:///home/conda/feedstock_root/build_artifacts/toolz_1657485559105/work
-torch==2.0.0
-torchaudio==2.0.1
-torchsummary==1.5.1
-torchvision==0.15.1
-tornado==6.3.1
-tqdm @ file:///home/conda/feedstock_root/build_artifacts/tqdm_1677948868469/work
-traitlets==5.9.0
-trimesh==3.21.5
-triton==2.0.0
-typing_extensions @ file:///home/conda/feedstock_root/build_artifacts/typing_extensions_1685704949284/work
-tzdata @ file:///home/conda/feedstock_root/build_artifacts/python-tzdata_1680081134351/work
-umap-learn==0.5.3
-unicodedata2 @ file:///home/conda/feedstock_root/build_artifacts/unicodedata2_1667239485250/work
-urllib3==1.26.16
-urllib3-secure-extra==0.1.0
-vit-pytorch==1.7.0
-wcwidth==0.2.6
-Werkzeug==2.3.3
-widgetsnbextension==4.0.7
-xxhash==3.2.0
-zfpc==0.1.2
-zfpy==1.0.0
-zipp @ file:///home/conda/feedstock_root/build_artifacts/zipp_1677313463193/work
-zmesh==1.0.0
-zope.event==4.6
-zope.interface==6.0
-zstandard==0.21.0
+# Affinity inference only. Install the CPU/CUDA PyTorch build for your platform first.
+# FRMC postprocessing: environment-postprocess.yml. Training: requirements-training.txt.
+torch>=2.6,<3
+numpy>=1.24,<3
+tifffile>=2023.7.10
diff --git a/tests/test_inference.py b/tests/test_inference.py
new file mode 100644
index 0000000..2d7c8c3
--- /dev/null
+++ b/tests/test_inference.py
@@ -0,0 +1,164 @@
+"""CPU-only correctness tests; no external data or downloaded weights needed."""
+
+import importlib.util
+import json
+from pathlib import Path
+import subprocess
+import sys
+import tempfile
+import unittest
+from unittest import mock
+from types import ModuleType
+
+import numpy as np
+import tifffile
+import torch
+
+
+ROOT = Path(__file__).resolve().parents[1]
+SPEC = importlib.util.spec_from_file_location("segneuron_inference", ROOT / "Train_and_Inference" / "inference.py")
+inference = importlib.util.module_from_spec(SPEC)
+SPEC.loader.exec_module(inference)
+
+
+class CoordinateModel(torch.nn.Module):
+ """Different channels expose coordinate transposition, normalization and blending errors."""
+ def __init__(self):
+ super().__init__()
+ self.register_buffer("fixture", torch.tensor(0))
+
+ def forward(self, x):
+ assert not self.training
+ assert torch.is_inference_mode_enabled()
+ return torch.cat((x, 1 - x, x / 2), dim=1), x * 0.75
+
+
+class InferenceTests(unittest.TestCase):
+ @classmethod
+ def setUpClass(cls):
+ cls.previous_threads = torch.get_num_threads()
+ torch.set_num_threads(1)
+
+ @classmethod
+ def tearDownClass(cls):
+ torch.set_num_threads(cls.previous_threads)
+
+ def test_reconstructs_small_odd_nonsquare_and_singleton_volumes(self):
+ for shape in ((1, 1, 1), (18, 31, 71), (21, 131, 197), (20, 128, 128)):
+ with self.subTest(shape=shape):
+ raw = np.random.default_rng(123).integers(0, 256, size=shape, dtype=np.uint8)
+ actual, boundary = inference.infer_volume(CoordinateModel(), raw)
+ normalized = raw.astype(np.float32) / 255
+ self.assertEqual(actual.shape, (3,) + shape)
+ self.assertEqual(boundary.shape, shape)
+ self.assertEqual(actual.dtype, np.float32)
+ self.assertTrue(np.isfinite(actual).all())
+ # Overlap sums use float32, as in the original Gaussian blend.
+ np.testing.assert_allclose(actual, np.stack((normalized, 1 - normalized, normalized / 2)), atol=1e-6)
+ np.testing.assert_allclose(boundary, normalized * 0.75, atol=1e-6)
+ self.assertGreaterEqual(actual.min(), 0)
+ self.assertLessEqual(actual.max(), 1)
+
+ def test_tiff_and_npy_load_identically(self):
+ with tempfile.TemporaryDirectory() as directory:
+ raw = np.arange(3 * 5 * 7, dtype=np.uint8).reshape(3, 5, 7)
+ for suffix in (".npy", ".tif"):
+ path = Path(directory) / ("raw" + suffix)
+ if suffix == ".npy":
+ np.save(path, raw)
+ else:
+ tifffile.imwrite(path, raw, photometric="minisblack")
+ np.testing.assert_array_equal(inference.load_volume(path), raw)
+
+ def test_rejects_invalid_raw_arrays(self):
+ for raw in (np.zeros((4, 4), np.uint8), np.zeros((0, 4, 4), np.uint8), np.zeros((4, 4, 4), np.uint16)):
+ with self.subTest(shape=raw.shape, dtype=raw.dtype):
+ with self.assertRaises(ValueError):
+ inference.infer_volume(CoordinateModel(), raw)
+
+ def test_rejects_rgb_and_time_series_tiff(self):
+ with tempfile.TemporaryDirectory() as directory:
+ path = Path(directory) / "raw.tif"
+ tifffile.imwrite(path, np.zeros((8, 9, 3), np.uint8), photometric="rgb")
+ with self.assertRaisesRegex(ValueError, "grayscale"):
+ inference.load_volume(path)
+ tifffile.imwrite(path, np.zeros((4, 8, 9), np.uint8), photometric="minisblack", metadata={"axes": "TYX"})
+ with self.assertRaisesRegex(ValueError, "grayscale"):
+ inference.load_volume(path)
+ tifffile.imwrite(path, np.zeros((4, 8, 9), np.uint8), photometric="minisblack", metadata={"axes": "YZX"})
+ with self.assertRaisesRegex(ValueError, "grayscale"):
+ inference.load_volume(path)
+
+ def test_checkpoint_formats_and_dataparallel_prefix(self):
+ source = torch.nn.Conv3d(1, 1, 1)
+ state = source.state_dict()
+ variants = [state, {"model_weights": {"module." + k: v for k, v in state.items()}}, {"state_dict": state}]
+ with tempfile.TemporaryDirectory() as directory:
+ path = Path(directory) / "model.pt"
+ for checkpoint in variants:
+ torch.save(checkpoint, path)
+ target = torch.nn.Conv3d(1, 1, 1)
+ inference.load_checkpoint(target, path)
+ for name, tensor in target.state_dict().items():
+ torch.testing.assert_close(tensor, state[name])
+
+ def test_checkpoint_mismatches_and_collisions_are_rejected(self):
+ with tempfile.TemporaryDirectory() as directory:
+ path = Path(directory) / "model.pt"
+ model = torch.nn.Conv3d(1, 1, 1)
+ torch.save({"weight": torch.zeros(1)}, path)
+ with self.assertRaises(RuntimeError):
+ inference.load_checkpoint(model, path)
+ torch.save({"weight": model.weight, "module.weight": model.weight}, path)
+ with self.assertRaisesRegex(ValueError, "duplicate"):
+ inference.load_checkpoint(model, path)
+
+ def test_nonfinite_model_output_is_rejected(self):
+ class InvalidModel(CoordinateModel):
+ def forward(self, x):
+ affinity, boundary = super().forward(x)
+ return affinity * float("nan"), boundary
+ with self.assertRaisesRegex(ValueError, "finite probabilities"):
+ inference.infer_volume(InvalidModel(), np.zeros((1, 2, 3), np.uint8))
+
+ def test_cli_help_does_not_require_site_packages(self):
+ result = subprocess.run([sys.executable, "-S", str(Path(inference.__file__)), "--help"], capture_output=True, text=True)
+ self.assertEqual(result.returncode, 0, result.stderr)
+ self.assertIn("--checkpoint", result.stdout)
+ self.assertIn("--device", result.stdout)
+
+ def test_cli_refuses_existing_output_before_loading_model(self):
+ with tempfile.TemporaryDirectory() as directory:
+ output = Path(directory)
+ raw, checkpoint = output / "raw.npy", output / "weights.pt"
+ np.save(raw, np.zeros((1, 2, 3), np.uint8))
+ torch.save({}, checkpoint)
+ with self.assertRaises(SystemExit) as caught:
+ inference.main(["--input", str(raw), "--checkpoint", str(checkpoint), "--output-dir", str(output)])
+ self.assertEqual(caught.exception.code, 2)
+ self.assertFalse((output / "inference.json").exists())
+
+ def test_cli_writes_complete_outputs_and_provenance(self):
+ # Exercise actual I/O and checkpoint loading with a small known model.
+ fake_package = ModuleType("model")
+ fake_module = ModuleType("model.Mnet")
+ fake_module.MNet = lambda *args, **kwargs: CoordinateModel()
+ with tempfile.TemporaryDirectory() as directory:
+ root = Path(directory)
+ raw, checkpoint, output = root / "raw.npy", root / "weights.pt", root / "output"
+ np.save(raw, np.full((1, 7, 11), 127, np.uint8))
+ torch.save({"model_weights": CoordinateModel().state_dict()}, checkpoint)
+ with mock.patch.dict(sys.modules, {"model": fake_package, "model.Mnet": fake_module}):
+ inference.main(["--input", str(raw), "--checkpoint", str(checkpoint), "--output-dir", str(output)])
+ affinity = np.load(output / "affinities.npy")
+ boundary = tifffile.imread(output / "boundaries.tif")
+ metadata = json.loads((output / "inference.json").read_text())
+ self.assertEqual(affinity.shape, (3, 1, 7, 11))
+ self.assertEqual(boundary.shape, (1, 7, 11))
+ self.assertEqual(metadata["input"]["sha256"], inference.sha256(raw))
+ self.assertEqual(metadata["checkpoint"]["sha256"], inference.sha256(checkpoint))
+ self.assertEqual(metadata["affinities"]["offsets_zyx"], [[-1, 0, 0], [0, -1, 0], [0, 0, -1]])
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/test_postprocess.py b/tests/test_postprocess.py
new file mode 100644
index 0000000..23edd1e
--- /dev/null
+++ b/tests/test_postprocess.py
@@ -0,0 +1,206 @@
+"""CLI/input regression tests plus opt-in-required tests of the real ELF backend."""
+
+import importlib.util
+import json
+import os
+from pathlib import Path
+import subprocess
+import sys
+import tempfile
+from types import SimpleNamespace
+import unittest
+from unittest.mock import patch
+
+import numpy as np
+import tifffile
+
+from Postprocess import FRMC_post as post
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+class PostprocessTests(unittest.TestCase):
+ def test_tiff_axes_are_not_silently_reinterpreted(self):
+ with tempfile.TemporaryDirectory() as temporary:
+ path = Path(temporary) / "head.tif"
+ tifffile.imwrite(path, np.zeros((4, 8, 9), dtype=np.float32),
+ photometric="minisblack", metadata={"axes": "YZX"})
+ with self.assertRaisesRegex(ValueError, "ZYX"):
+ post._load_volume(path)
+
+ def test_help_does_not_import_elf(self):
+ code = (
+ "import sys; sys.modules['elf'] = None; "
+ "from Postprocess.FRMC_post import main; main(['--help'])"
+ )
+ result = subprocess.run([sys.executable, "-c", code], cwd=ROOT, capture_output=True, text=True)
+ self.assertEqual(result.returncode, 0, result.stderr)
+ self.assertIn("--ground-truth", result.stdout)
+
+ def test_invalid_input_is_rejected_before_backend_load(self):
+ valid = np.ones((3, 2, 8, 8), dtype=np.float32)
+ invalid = [np.ones((2, 2, 8, 8)), np.empty((3, 0, 8, 8)),
+ np.full_like(valid, np.nan), np.full_like(valid, 1.01),
+ np.full_like(valid, -0.01), valid.astype(complex)]
+ with patch.object(post, "_load_elf", side_effect=AssertionError("backend must not load")):
+ for values in invalid:
+ with self.subTest(shape=values.shape), self.assertRaises(ValueError):
+ post.post_mc(values)
+ for beta in [0, 1, -1, np.nan, np.inf, True, None, [0.25]]:
+ with self.subTest(beta=beta), self.assertRaises(ValueError):
+ post.post_mc(valid, beta)
+
+ def test_fragments_are_disjoint_across_slices_and_solver_zero_is_preserved(self):
+ fragment_slice = np.array([[1, 1, 9], [1, 9, 9]], dtype=np.uint64)
+ captured = {}
+
+ def compute_rag(watershed):
+ captured["watershed"] = watershed.copy()
+ return SimpleNamespace(numberOfEdges=3, numberOfNodes=4)
+
+ def project(rag, labels):
+ return labels[captured["watershed"]]
+
+ feats = SimpleNamespace(
+ compute_rag=compute_rag,
+ compute_affinity_features=lambda *a: np.full((3, 1), 0.5),
+ compute_boundary_mean_and_length=lambda *a: np.ones((3, 2)),
+ project_node_labels_to_pixels=project,
+ )
+ mc = SimpleNamespace(
+ transform_probabilities_to_costs=lambda costs, **kw: costs,
+ multicut_kernighan_lin=lambda *a: np.array([0, 7, 0, 7]),
+ )
+ ws = SimpleNamespace(distance_transform_watershed=lambda *a, **kw: (fragment_slice.copy(), 9))
+ with patch.object(post, "_load_elf", return_value=(feats, mc, ws)):
+ result = post.post_mc(np.full((3, 2, 2, 3), 0.8))
+ self.assertEqual(set(np.unique(captured["watershed"][0])), {0, 1})
+ self.assertEqual(set(np.unique(captured["watershed"][1])), {2, 3})
+ self.assertEqual(result.dtype, np.uint32)
+ np.testing.assert_array_equal(result[0], result[1])
+ np.testing.assert_array_equal(np.unique(result), [1, 2])
+
+ def test_unassigned_watershed_is_an_error(self):
+ ws = SimpleNamespace(distance_transform_watershed=lambda *a, **kw: (np.zeros((8, 8), dtype=np.uint64), 0))
+ with patch.object(post, "_load_elf", return_value=(None, None, ws)):
+ with self.assertRaisesRegex(RuntimeError, "unassigned"):
+ post.post_mc(np.ones((3, 1, 8, 8)))
+
+ def test_cli_fusion_output_and_optional_metrics(self):
+ with tempfile.TemporaryDirectory() as temporary:
+ directory = Path(temporary)
+ affs = np.full((3, 2, 8, 8), 0.9, dtype=np.float32)
+ foreground = np.full((2, 8, 8), 0.7, dtype=np.float32)
+ labels = np.ones(foreground.shape, dtype=np.uint32)
+ labels[:, :, 4:] = 2
+ np.save(directory / "affs.npy", affs)
+ tifffile.imwrite(directory / "head.tif", foreground, photometric="minisblack")
+ tifffile.imwrite(directory / "gt.tif", labels, photometric="minisblack")
+ base = ["--affinities", str(directory / "affs.npy"), "--boundaries", str(directory / "head.tif")]
+ for extension, with_gt in [("tif", False), ("npy", True)]:
+ output = directory / f"labels.{extension}"
+ args = base + ["--output", str(output)]
+ if with_gt:
+ args += ["--ground-truth", str(directory / "gt.tif")]
+ with patch.object(post, "post_mc", return_value=labels) as solver, patch("builtins.print") as printed:
+ self.assertEqual(post.main(args), 0)
+ np.testing.assert_array_equal(solver.call_args.args[0], np.minimum(affs, foreground[None]))
+ self.assertEqual(solver.call_args.args[1], 0.25)
+ np.testing.assert_array_equal(post._load_volume(output), labels)
+ summary = json.loads(printed.call_args.args[0])
+ self.assertEqual("metrics" in summary, with_gt)
+ if with_gt:
+ self.assertEqual(summary["metrics"]["arand"], 0)
+ self.assertEqual(summary["metrics"]["voi"], 0)
+ original_bytes = output.read_bytes()
+ with patch.object(post, "post_mc", side_effect=AssertionError("must fail first")):
+ with self.assertRaises(SystemExit):
+ post.main(args)
+ self.assertEqual(output.read_bytes(), original_bytes)
+
+ def test_backend_failure_does_not_leave_output(self):
+ with tempfile.TemporaryDirectory() as temporary:
+ directory = Path(temporary)
+ np.save(directory / "affs.npy", np.full((3, 2, 8, 8), 0.8))
+ tifffile.imwrite(directory / "head.tif", np.full((2, 8, 8), 0.8), photometric="minisblack")
+ output = directory / "out.npy"
+ with patch.object(post, "_load_elf", side_effect=RuntimeError("backend unavailable")):
+ with self.assertRaises(SystemExit):
+ post.main(["--affinities", str(directory / "affs.npy"), "--boundaries",
+ str(directory / "head.tif"), "--output", str(output)])
+ self.assertFalse(output.exists())
+
+ def test_ground_truth_validation_and_sparse_labels(self):
+ for invalid in [np.zeros((2, 3), dtype=np.uint32), np.full((2, 3), -1), np.ones((2, 3), dtype=float)]:
+ with self.assertRaises(ValueError):
+ post._validate_ground_truth(invalid, (2, 3))
+ with self.assertRaisesRegex(ValueError, "shape"):
+ post._validate_ground_truth(np.ones((2, 3), dtype=np.uint32), (3, 2))
+ compact = post._validate_ground_truth(np.array([[0, 2**40], [0, 9]], dtype=np.uint64), (2, 2))
+ np.testing.assert_array_equal(compact, [[0, 2], [0, 1]])
+
+ def test_exclusive_write_preserves_existing_file(self):
+ with tempfile.TemporaryDirectory() as temporary:
+ output = Path(temporary) / "labels.npy"
+ post._write_labels(output, np.ones((2, 3, 4), dtype=np.uint32))
+ before = output.read_bytes()
+ with self.assertRaises(FileExistsError):
+ post._write_labels(output, np.zeros((2, 3, 4), dtype=np.uint32))
+ self.assertEqual(output.read_bytes(), before)
+
+
+class RealElfTests(unittest.TestCase):
+ @classmethod
+ def setUpClass(cls):
+ try:
+ post._load_elf()
+ except RuntimeError as exc:
+ if os.environ.get("SEGNEURON_REQUIRE_ELF") == "1":
+ raise
+ raise unittest.SkipTest(str(exc)) from exc
+
+ @staticmethod
+ def affinities():
+ # Two large regions separated by a strong x boundary, linked across z.
+ affs = np.full((3, 3, 48, 64), 0.98, dtype=np.float32)
+ affs[2, :, :, 31:34] = 0.02
+ return affs
+
+ def test_real_multicut_preserves_original_partition(self):
+ specification = importlib.util.spec_from_file_location("legacy_postprocess", ROOT / "legacy/Postprocess/FRMC_post.py")
+ legacy = importlib.util.module_from_spec(specification)
+ specification.loader.exec_module(legacy)
+ affs = self.affinities()
+ before = affs.copy()
+ for beta in (0.1, 0.25, 0.5, 0.75):
+ with self.subTest(beta=beta):
+ original = legacy.post_mc(affs, beta)
+ result = post.post_mc(affs, beta)
+ pairs = np.unique(np.stack([original.ravel(), result.ravel()], axis=1), axis=0)
+ self.assertEqual(len(pairs), len(np.unique(original)))
+ self.assertEqual(len(pairs), len(np.unique(result)))
+ self.assertEqual(result.shape, affs.shape[1:])
+ self.assertEqual(result.dtype, np.uint32)
+ self.assertGreater(result.min(), 0)
+ np.testing.assert_array_equal(affs, before)
+
+ def test_real_cli_without_ground_truth(self):
+ with tempfile.TemporaryDirectory() as temporary:
+ directory = Path(temporary)
+ affinities = self.affinities()
+ np.save(directory / "affs.npy", affinities)
+ tifffile.imwrite(directory / "head.tif", np.ones(affinities.shape[1:], dtype=np.float32), photometric="minisblack")
+ output = directory / "labels.tif"
+ with patch("builtins.print") as printed:
+ post.main(["--affinities", str(directory / "affs.npy"), "--boundaries",
+ str(directory / "head.tif"), "--output", str(output)])
+ labels = tifffile.imread(output)
+ self.assertEqual(labels.dtype, np.uint32)
+ self.assertEqual(labels.shape, affinities.shape[1:])
+ self.assertGreater(labels.min(), 0)
+ self.assertNotIn("metrics", json.loads(printed.call_args.args[0]))
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/test_training.py b/tests/test_training.py
new file mode 100644
index 0000000..8b3a73a
--- /dev/null
+++ b/tests/test_training.py
@@ -0,0 +1,89 @@
+"""Exercise research training plumbing without the full EMNeuron corpus or a GPU."""
+import importlib.util
+from pathlib import Path
+import subprocess
+import sys
+import tempfile
+import unittest
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def smoke(kind):
+ import random
+ from unittest.mock import Mock, patch
+ from addict import Dict
+ import numpy as np
+ import torch
+ import yaml
+
+ folder, name = (('Pretrain', 'pretrain') if kind == 'pretrain'
+ else ('Train_and_Inference', 'supervised_train'))
+ sys.path.insert(0, str(ROOT / folder))
+ module = __import__(name)
+ with (ROOT / folder / 'config' / 'SegNeuron.yaml').open() as handle:
+ cfg = Dict(yaml.safe_load(handle))
+ assert cfg.MODEL.model_type != 'mala'
+
+ class TinyModel(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.bias = torch.nn.Parameter(torch.tensor(0.0))
+
+ def forward(self, inputs):
+ output = torch.sigmoid(inputs + self.bias)
+ return (output if kind == 'pretrain' else output.repeat(1, 3, 1, 1, 1), output)
+
+ inputs = torch.zeros((1, 1, 2, 8, 8))
+ target = torch.ones_like(inputs)
+ batch = ((inputs, target, target) if kind == 'pretrain' else
+ (inputs, target.repeat(1, 3, 1, 1, 1), target, target))
+ provider = Mock()
+ provider.next.side_effect = [batch, batch]
+ with patch.object(module, 'Provider', return_value=provider):
+ assert module.load_dataset(cfg) is provider
+
+ with tempfile.TemporaryDirectory() as directory:
+ cfg.record_path = cfg.cache_path = cfg.save_path = directory
+ cfg.TRAIN.total_iters = 2
+ cfg.TRAIN.display_freq = cfg.TRAIN.valid_freq = 1
+ cfg.TRAIN.save_freq = 2
+ cfg.TRAIN.base_lr = cfg.TRAIN.end_lr = 0.01
+ module.cfg = cfg
+ model = TinyModel()
+ optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+ module.loop(cfg, provider, model, optimizer, 0, Mock())
+ assert provider.next.call_count == 2
+ assert float(model.bias.detach()) > 0
+ assert (Path(directory) / '000001.png').is_file()
+ checkpoint = torch.load(Path(directory) / 'model-000002.ckpt', weights_only=True)
+ assert checkpoint['current_iter'] == 2
+
+ if kind == 'pretrain':
+ from pretrain_provider import Train
+ dataset = Train.__new__(Train)
+ dataset.dataset = [np.zeros((20, 128, 128), dtype=np.uint8)]
+ dataset.crop_from_origin = [20, 128, 128]
+ dataset.simple_aug = lambda arrays: arrays
+ random.seed(0)
+ np.random.seed(0)
+ outputs = dataset[0]
+ assert all(value.shape == (1, 20, 128, 128) for value in outputs)
+ assert all(np.isfinite(value).all() for value in outputs)
+
+
+class TrainingTests(unittest.TestCase):
+ @unittest.skipUnless(importlib.util.find_spec('addict'), 'Install requirements-training.txt')
+ def test_research_entrypoints_load_optimize_render_and_save(self):
+ for kind in ('pretrain', 'supervised'):
+ with self.subTest(kind=kind):
+ result = subprocess.run([sys.executable, str(Path(__file__).resolve()), '--smoke', kind],
+ capture_output=True, text=True, timeout=90)
+ self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
+
+
+if __name__ == '__main__':
+ if len(sys.argv) == 3 and sys.argv[1] == '--smoke':
+ smoke(sys.argv[2])
+ else:
+ unittest.main()