OSCR

Unsupervised deep learning enables blur-free resolution enhancement in two-photon microscopy.

Code ↔ Paper

5 matches between paragraphs of the paper and lines of its authors' code, computed by the harvester (lexical-v1). Click a colored paragraph or line to see its counterpart.

The 5 matches
  1. [1] § STAR★Methods › Methods details › Hyperparameters for training ↔ model.py, lines 496–553 · score 0.64 · Gibson Lanni model, tg, wavelength, NA
  2. [2] § STAR★Methods › Methods details › Blur generation model architecture ↔ model.py, lines 403–471 · score 0.60 · Gibson Lanni model, neural implicit, PSF, Blur
  3. [3] § STAR★Methods › Methods details › Network architecture ↔ model.py, lines 403–471 · score 0.60 · Gibson Lanni model, neural implicit, PSF, trained
  4. [4] § STAR★Methods › Quantification and statistical analysis › Hyperparameter sensitivity analysis ↔ figures/paper_figures_suppli.ipynb, lines 76–199 · score 0.56 · VQ loss, EWC loss, PSF loss, MRF
  5. [5] § STAR★Methods › Methods details › Optimization ↔ model.py, lines 328–400 · score 0.51 · neural implicit PSF, RProp, batch, trained

Paper

Loaded from Europe PMC by your browser, not stored by OSCR: doi.org · Europe PMC

The paper is loaded when this pane is shown.

The authors' code

Python · 775 lines · 31 KB · no license · 4 matches

  1. import numpy as np
  2. import scipy
  3. import torch
  4. import torch.nn as nn
  5. import torch.nn.functional as F
  6. import torch.distributions as dist
  7. from torch.utils.checkpoint import checkpoint
  8. import torch.special as S
  9. from scipy.stats import lognorm
  10. from fft_conv_pytorch import fft_conv
  11. import matplotlib.pyplot as plt
  12. import time
  14. class JNetBlock0(nn.Module):
  15. def __init__(self, in_channels, out_channels):
  16. super().__init__()
  17. self.conv = nn.Conv3d(in_channels = in_channels ,
  18. out_channels = out_channels,
  19. kernel_size = 7 ,
  20. padding = 'same' ,
  21. padding_mode = 'replicate' ,)
  23. def forward(self, x):
  24. x = self.conv(x)
  25. return x
  28. class JNetBlock(nn.Module):
  29. def __init__(self, in_channels, hidden_channels, dropout):
  30. super().__init__()
  31. self.bn1 = nn.BatchNorm3d(num_features = in_channels)
  32. self.relu1 = nn.ReLU(inplace=True)
  33. self.conv1 = nn.Conv3d(in_channels = in_channels ,
  34. out_channels = hidden_channels,
  35. kernel_size = 3 ,
  36. padding = 'same' ,
  37. padding_mode = 'replicate' ,)
  39. self.bn2 = nn.BatchNorm3d(num_features = hidden_channels)
  40. self.relu2 = nn.ReLU(inplace=True)
  41. self.dropout1 = nn.Dropout(p = dropout)
  42. self.conv2 = nn.Conv3d(in_channels = hidden_channels,
  43. out_channels = in_channels ,
  44. kernel_size = 3 ,
  45. padding = 'same' ,
  46. padding_mode = 'replicate' ,)
  48. def forward(self, x):
  49. d = self.bn1(x)
  50. d = self.relu1(d)
  51. d = self.conv1(d)
  52. d = self.bn2(d)
  53. d = self.relu2(d)
  54. d = self.dropout1(d)
  55. d = self.conv2(d)
  56. x = x + d
  57. return x
  60. class JNetBlockN(nn.Module):
  61. def __init__(self, in_channels, out_channels):
  62. super().__init__()
  63. self.conv = nn.Conv3d(in_channels = in_channels ,
  64. out_channels = out_channels,
  65. kernel_size = 3 ,
  66. padding = 'same' ,
  67. padding_mode = 'replicate' ,)
  68. self.sigm = nn.Sigmoid()
  70. def forward(self, x):
  71. x = self.conv(x)
  72. x = self.sigm(x)
  73. return x
  76. class JNetPooling(nn.Module):
  77. def __init__(self, in_channels, out_channels):
  78. super().__init__()
  79. self.maxpool = nn.MaxPool3d(kernel_size = 2)
  80. self.conv = nn.Conv3d(in_channels = in_channels ,
  81. out_channels = out_channels ,
  82. kernel_size = 1 ,
  83. padding = 'same' ,
  84. padding_mode = 'replicate' ,)
  85. self.relu = nn.ReLU(inplace=True)
  87. def forward(self, x):
  88. x = self.maxpool(x)
  89. x = self.conv(x)
  90. x = self.relu(x)
  91. return x
  94. class JNetUnpooling(nn.Module):
  95. def __init__(self, in_channels, out_channels):
  96. super().__init__()
  97. self.upsample = nn.Upsample(scale_factor = 2 ,
  98. mode = 'trilinear' ,)
  99. self.conv = nn.Conv3d(in_channels = in_channels ,
  100. out_channels = out_channels ,
  101. kernel_size = 1 ,
  102. padding = 'same' ,
  103. padding_mode = 'replicate' ,)
  104. self.relu = nn.ReLU(inplace=True)
  106. def forward(self, x):
  107. x = self.upsample(x)
  108. x = self.conv(x)
  109. x = self.relu(x)
  110. return x
  113. class JNetUpsample(nn.Module):
  114. def __init__(self, scale_factor):
  115. super().__init__()
  116. self.upsample = nn.Upsample(scale_factor = scale_factor ,
  117. mode = 'trilinear' ,)
  119. def forward(self, x):
  120. return self.upsample(x)
  123. class SuperResolutionBlock(nn.Module):
  124. def __init__(self, scale_factor, in_channels, nblocks, dropout):
  125. super().__init__()
  126. self.upsample = nn.Upsample(scale_factor = scale_factor ,
  127. mode = 'trilinear' ,)
  128. self.post = nn.ModuleList([JNetBlock(in_channels = in_channels ,
  129. hidden_channels = in_channels ,
  130. dropout = dropout ,
  131. ) for _ in range(nblocks)])
  132. def forward(self, x):
  133. x = self.upsample(x)
  134. for f in self.post:
  135. x = f(x)
  136. return x
  139. class CrossAttentionBlock(nn.Module):
  140. """
  141. ### Transformer Layer
  142. """
  144. def __init__(self, channels: int, n_heads: int, d_cond: int):
  145. """
  146. :param d_model: is the input embedding size
  147. :param n_heads: is the number of attention heads
  148. :param d_head: is the size of a attention head
  149. :param d_cond: is the size of the conditional embeddings
  150. """
  151. super().__init__()
  152. self.attn = CrossAttention(d_model = channels,
  153. d_cond = d_cond,
  154. n_heads = n_heads,
  155. d_head = channels // n_heads,)
  156. self.norm = nn.LayerNorm(normalized_shape = channels,)
  158. def forward(self, x: torch.Tensor):
  159. """
  160. :param x: are the input embeddings of shape `[batch_size, height * width, d_model]`
  161. :param cond: is the conditional embeddings of shape `[batch_size, n_cond, d_cond]`
  162. """
  163. b, c, d, h, w = x.shape
  164. x = x.permute(0, 2, 3, 4, 1).view(b, d * h * w, c)
  165. x = self.attn(self.norm(x)) + x
  166. x = x.view(b, d, h, w, c).permute(0, 4, 1, 2, 3)
  167. return x
  170. class CrossAttention(nn.Module):
  171. """
  172. ### Cross Attention Layer
  173. This falls-back to self-attention when conditional embeddings are not specified.
  174. """
  176. def __init__(self, d_model: int, d_cond: int, n_heads: int, d_head: int, is_inplace: bool = True):
  177. """
  178. :param d_model: is the input embedding size
  179. :param n_heads: is the number of attention heads
  180. :param d_head: is the size of a attention head
  181. :param d_cond: is the size of the conditional embeddings
  182. :param is_inplace: specifies whether to perform the attention softmax computation inplace to
  183. save memory
  184. """
  185. super().__init__()
  187. self.is_inplace = is_inplace
  188. self.n_heads = n_heads
  189. self.d_head = d_head
  191. # Attention scaling factor
  192. self.scale = d_head ** -0.5
  193. # Query, key and value mappings
  194. d_attn = d_head * n_heads
  195. self.to_q = nn.Linear(in_features = d_model ,
  196. out_features = d_attn ,
  197. bias = False ,)
  199. self.to_k = nn.Linear(in_features = d_cond ,
  200. out_features = d_attn ,
  201. bias = False ,)
  203. self.to_v = nn.Linear(in_features = d_cond ,
  204. out_features = d_attn ,
  205. bias = False ,)
  206. # Final linear layer
  207. self.to_out = nn.Sequential(
  208. nn.Linear(in_features = d_attn ,
  209. out_features = d_model,),
  210. )
  212. def forward(self, x: torch.Tensor, cond=None):
  213. """
  214. :param x: are the input embeddings of shape `[batch_size, height * width, d_model]`
  215. :param cond: is the conditional embeddings of shape `[batch_size, n_cond, d_cond]`
  216. """
  218. has_cond = cond is not None
  219. if not has_cond:
  220. cond = x
  221. q = self.to_q(x)
  222. k = self.to_k(cond)
  223. v = self.to_v(cond)
  225. return self.normal_attention(q, k, v)
  227. def normal_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
  228. """
  229. #### Normal Attention
  230. :param q: `[batch_size, seq, d_attn]`
  231. :param k: `[batch_size, seq, d_attn]`
  232. :param v: `[batch_size, seq, d_attn]`
  233. """
  235. q = q.view(*q.shape[:2], self.n_heads, -1)
  236. k = k.view(*k.shape[:2], self.n_heads, -1)
  237. v = v.view(*v.shape[:2], self.n_heads, -1)
  238. attn = torch.einsum('bihd,bjhd->bhij', q, k) * self.scale
  240. if self.is_inplace:
  241. half = attn.shape[0] // 2
  242. attn[half:] = attn[half:].softmax(dim=-1)
  243. attn[:half] = attn[:half].softmax(dim=-1)
  244. else:
  245. attn = attn.softmax(dim=-1)
  247. out = torch.einsum('bhij,bjhd->bihd', attn, v)
  248. # Reshape to `[batch_size, height * width * depth, n_heads * d_head]`
  249. out = out.reshape(*out.shape[:2], -1)
  250. # Map to `[batch_size, height * width, d_model]` with a linear layer
  251. return self.to_out(out)
  254. class VectorQuantizer(nn.Module):
  255. def __init__(self, threshold, device):
  256. super().__init__()
  257. self.t = threshold
  258. self.device = device
  259. def forward(self, x):
  260. x_quantized = (x >= self.t).to(self.device).float()
  261. x_quantized = x + (x_quantized - x).detach()
  262. quantize_loss = F.mse_loss(x_quantized.detach(), x)
  263. return x_quantized, quantize_loss
  266. class JNetLayer(nn.Module):
  267. def __init__(self, in_channels, hidden_channels_list,
  268. attn_list, nblocks, dropout):
  269. super().__init__()
  270. is_attn = attn_list.pop(0)
  271. hidden_channels = hidden_channels_list.pop(0)
  272. self.hidden_channels = hidden_channels
  273. self.pool = JNetPooling(in_channels = in_channels ,
  274. out_channels = hidden_channels,)
  275. self.conv = nn.Conv3d(in_channels = hidden_channels,
  276. out_channels = hidden_channels,
  277. kernel_size = 1 ,
  278. padding = 'same' ,
  279. padding_mode = 'replicate' ,)
  280. self.prev = nn.ModuleList([JNetBlock(in_channels = hidden_channels,
  281. hidden_channels = hidden_channels,
  282. dropout = dropout ,
  283. ) for _ in range(nblocks)])
  284. self.mid = JNetLayer(in_channels = hidden_channels ,
  285. hidden_channels_list = hidden_channels_list ,
  286. attn_list = attn_list ,
  287. nblocks = nblocks ,
  288. dropout = dropout ,
  289. ) if hidden_channels_list else nn.Identity()
  290. self.attn = CrossAttentionBlock(channels = hidden_channels ,
  291. n_heads = 8 ,
  292. d_cond = hidden_channels ,)\
  293. if is_attn else nn.Identity()
  294. self.post = nn.ModuleList([JNetBlock(in_channels = hidden_channels,
  295. hidden_channels = hidden_channels,
  296. dropout = dropout ,
  297. ) for _ in range(nblocks)])
  298. self.unpool = JNetUnpooling(in_channels = hidden_channels,
  299. out_channels = in_channels ,)
  301. def forward(self, x):
  302. d = self.pool(x)
  303. d = self.conv(d) # checkpoint
  304. for f in self.prev:
  305. d = f(d)
  306. d = self.mid(d)
  307. d = self.attn(d)
  308. for f in self.post:
  309. d = f(d)
  310. d = self.unpool(d) # checkpoint
  311. x = x + d
  312. return x
  315. class Emission(nn.Module):
  316. def __init__(self):
  317. super().__init__()
  319. def sample(self, x, params):
  320. b = x.shape[0]
  321. pz0 = dist.LogNormal(loc = torch.tensor(float(params["mu_z"])).view( b,1,1,1,1).expand(*x.shape),
  322. scale = torch.tensor(float(params["sig_z"])).view(b,1,1,1,1).expand(*x.shape),)
  323. x = x * pz0.sample().to(x.device)
  324. x = torch.clip(x, min=0., max=1.)
  325. return x
  328. class NeuralImplicitPSF(nn.Module):
  329. def __init__(self, config, psf):
  330. super().__init__()
  331. mid = config["mid"]
  332. self.layers = nn.Sequential(
  333. nn.BatchNorm1d(2) ,
  334. nn.Linear(2, mid) ,
  335. nn.Sigmoid() ,
  336. nn.BatchNorm1d(mid),
  337. nn.Linear(mid, 1) ,
  338. nn.Sigmoid()
  339. )
  340. self.config = config
  341. self._gen_coord(psf)
  342. self._gen_label(psf)
  344. def _init_weights(self):
  345. for m in self.parameters():
  346. torch.nn.init.normal_(m)
  348. def forward(self, coords):
  349. return self.layers(coords)
  351. def trainer(self):
  352. self.to(self.config["device"])
  353. self._train()
  355. def _gen_coord(self, psf):
  356. self.psf_shape = psf.shape
  357. z, r = self.psf_shape
  358. zs = torch.linspace(-1, 1, steps=z).abs()
  359. rs = torch.linspace(-1, 1, steps=r)
  360. grid_z, grid_r = torch.meshgrid(zs, rs, indexing='ij')
  361. self.coord = torch.stack((
  362. grid_z.flatten(), grid_r.flatten()), -1).to(self.config["device"])
  364. def _gen_label(self, psf):
  365. self.label = psf.flatten()[:, None]
  367. def _train(self):
  368. label = self.label.to(self.config["device"])
  369. coord = self.coord.to(self.config["device"])
  370. loss_fn = eval(self.config["loss_fn"])
  371. loss = torch.tensor(1.)
  372. while loss.item() >= self.config["nipsf_loss_target"]:
  373. self._init_weights()
  374. count = 0
  375. label = label.to(self.config["device"])
  376. coord = coord.to(self.config["device"])
  377. optim = torch.optim.Rprop(self.parameters(), lr=self.config["lr"])
  378. averaged_model = torch.optim.swa_utils.AveragedModel(
  379. self,
  380. multi_avg_fn=torch.optim.swa_utils.get_ema_multi_avg_fn(0.999))
  381. sched = torch.optim.lr_scheduler.ExponentialLR(
  382. optim,
  383. gamma=0.999,)
  384. while loss.item() >= self.config["nipsf_loss_target"]:
  385. optim.zero_grad()
  386. loss = loss_fn(self(coord), label)
  387. loss.backward()
  388. optim.step()
  389. averaged_model.update_parameters(self)
  390. sched.step()
  391. count += 1
  392. if count == self.config["num_iter_psf_pretrain"]:
  393. print("Neural Implicit PSF is not good enough. " \
  394. "Initializing NeuriPSF and trying again...")
  395. break
  396. print(f"NeuriPSF train done with loss target"\
  397. f" {self.config['nipsf_loss_target']}.")
  399. def array(self):
  400. return self(self.coord).view(self.psf_shape)
  403. class Blur(nn.Module):
  404. def __init__(self, params):
  405. super().__init__()
  406. self.device = params["device"]
  407. self.use_fftconv = params["use_fftconv"]
  408. if params["blur_mode"] == "gaussian":
  409. psf_model = GaussianModel(params)
  410. elif params["blur_mode"] == "gibsonlanni":
  411. psf_model = GibsonLanniModel(params)
  412. else:
  413. raise(NotImplementedError(
  414. f'blur_mode {params["blur_mode"]} is not implemented. Try "gaussian" or "gibsonlanni".'))
  415. self.init_psf_rz = torch.tensor(psf_model.PSF_rz, requires_grad=False).float().to(self.device)
  416. self.neuripsf = NeuralImplicitPSF(params, self.init_psf_rz)
  417. # self.neuripsf.trainer(self.init_psf_rz) <- moved to train_runner.py
  418. self.psf_rz_s0 = self.init_psf_rz.shape[0]
  419. self.size_z = params["size_z"]
  420. xy = torch.meshgrid(torch.arange(params["size_y"]),
  421. torch.arange(params["size_x"]),
  422. indexing='ij')
  423. r = torch.tensor(psf_model.r)
  424. x0 = (params["size_x"] - 1) / 2
  425. y0 = (params["size_y"] - 1) / 2
  426. r_pixel = torch.sqrt((xy[1] - x0) ** 2 + (xy[0] - y0) ** 2) * params["res_lateral"]
  427. rs0, = r.shape
  428. self.rps0, self.rps1 = r_pixel.shape
  429. r_e = r[:, None, None].expand(rs0, self.rps0, self.rps1)
  430. r_pixel_e = r_pixel[None].expand(rs0, self.rps0, self.rps1)
  431. r_index = torch.argmin(torch.abs(r_e- r_pixel_e), dim=0)
  432. r_index_fe = r_index.flatten().expand(self.psf_rz_s0, -1)
  433. self.r_index_fe = r_index_fe.to(self.device)
  434. self.z_pad = int((params["size_z"] - params["res_axial"] // params["res_lateral"] + 1) // 2)
  435. self.x_pad = (params["size_x"]) // 2
  436. self.y_pad = (params["size_y"]) // 2
  437. self.stride = (params["scale"], 1, 1)
  439. def forward(self, x):
  440. psf_rz = self.neuripsf.array()
  441. l2_psf_rz = torch.mean((psf_rz - self.init_psf_rz) ** 2)
  442. psf = torch.gather(psf_rz, 1, self.r_index_fe)
  443. psf = psf.reshape(self.psf_rz_s0, self.rps0, self.rps1)
  444. psf = F.interpolate(
  445. input = psf[None, None, :] ,
  446. size = (self.size_z, self.rps0, self.rps1) ,
  447. mode = "nearest" ,
  448. )[0, 0]
  449. psf = psf / torch.sum(psf)#;print("sum: ", torch.sum(psf));print("max: ", torch.max(psf))#/ torch.sum(psf)
  450. if self.use_fftconv:
  451. _x = fft_conv(signal = x ,
  452. kernel = psf ,
  453. stride = self.stride ,
  454. padding = (self.z_pad, self.x_pad, self.y_pad,),
  455. )
  456. else:
  457. _x = F.conv3d(input = x ,
  458. weight = psf ,
  459. stride = self.stride ,
  460. padding = (self.z_pad, self.x_pad, self.y_pad,),
  461. )
  462. return {"out" : _x ,
  463. "psf_loss": l2_psf_rz}
  465. def show_psf_3d(self):
  466. # with torch.no_grad():
  467. psf_rz = self.neuripsf.array()
  468. psf = torch.gather(psf_rz, 1, self.r_index_fe)
  469. psf = psf / torch.sum(psf)
  470. psf = psf.reshape(self.psf_rz_s0, self.rps0, self.rps1)
  471. return psf
  474. class GaussianModel():
  475. def __init__(self, params):
  476. oversampling = 1 # Defines the upsampling ratio on the image space grid for computations
  477. size_x = params["size_x"]
  478. size_y = params["size_y"]
  479. size_z = params["size_z"] // params["scale"]
  480. bet_xy = params["bet_xy"]
  481. bet_z = params["bet_z" ]
  482. x0 = (size_x - 1) / 2
  483. y0 = (size_y - 1) / 2
  484. z0 = (size_z - 1) / 2
  485. res_lateral = params["res_lateral"]#0.05 # microns # # # # param # # # #
  486. max_radius = round(np.sqrt((size_x - x0) * (size_x - x0) + (size_y - y0) * (size_y - y0)))
  487. self.r = res_lateral * np.arange(0, oversampling * max_radius) / oversampling
  488. xy = np.meshgrid(np.arange(size_z), np.arange(max_radius), indexing="ij")
  489. distance = np.sqrt((xy[1] / bet_xy) ** 2 + ((xy[0] - z0) / bet_z) ** 2)
  490. self.PSF_rz = np.exp(- distance ** 2)
  492. def __call__(self):
  493. return self.PSF_rz
  496. class GibsonLanniModel():
  497. def __init__(self, params):
  498. size_x = params["size_x"]#256 # # # # param # # # #
  499. size_y = params["size_y"]#256 # # # # param # # # #
  500. size_z = params["size_z"] // params["scale"]#128 # # # # param # # # #
  502. # Precision control
  503. num_basis = 100 # Number of rescaled Bessels that approximate the phase function
  504. num_samples = 1000 # Number of pupil samples along radial direction
  505. oversampling = 1 # Defines the upsampling ratio on the image space grid for computations
  507. # Microscope parameters
  508. NA = params["NA"] #1.1 # # # # param # # # #
  509. wavelength = params["wavelength"]#0.910 # microns # # # # param # # # #
  510. M = params["M"] #25 # magnification # # # # param # # # #
  511. ns = params["ns"] #1.33 # specimen refractive index (RI)
  512. ng0 = params["ng0"]#1.5 # coverslip RI design value
  513. ng = params["ng"] #1.5 # coverslip RI experimental value
  514. ni0 = params["ni0"]#1.5 # immersion medium RI design value
  515. ni = params["ni"] #1.5 # immersion medium RI experimental value
  516. ti0 = params["ti0"]#150 # microns, working distance (immersion medium thickness) design value
  517. tg0 = params["tg0"]#170 # microns, coverslip thickness design value
  518. tg = params["tg"] #170 # microns, coverslip thickness experimental value
  519. res_lateral = params["res_lateral"]#0.05 # microns # # # # param # # # #
  520. res_axial = params["res_axial"]#0.5 # microns # # # # param # # # #
  521. pZ = params["pZ"] # 2 microns, particle distance from coverslip
  523. # Scaling factors for the Fourier-Bessel series expansion
  524. min_wavelength = 0.436 # microns
  525. scaling_factor = NA * (3 * np.arange(1, num_basis + 1) - 2) * min_wavelength / wavelength
  526. x0 = (size_x - 1) / 2
  527. y0 = (size_y - 1) / 2
  528. max_radius = round(np.sqrt((size_x - x0) * (size_x - x0) + (size_y - y0) * (size_y - y0)))
  529. r = res_lateral * np.arange(0, oversampling * max_radius) / oversampling
  530. self.r = r
  531. a = min([NA, ns, ni, ni0, ng, ng0]) / NA
  532. rho = np.linspace(0, a, num_samples)
  534. z = res_axial * np.arange(-size_z / 2, size_z /2) + res_axial / 2
  536. OPDs = pZ * np.sqrt(ns * ns - NA * NA * rho * rho) # OPD in the sample
  537. OPDi = (z.reshape(-1,1) + ti0) * np.sqrt(ni * ni - NA * NA * rho * rho) - ti0 * np.sqrt(ni0 * ni0 - NA * NA * rho * rho) # OPD in the immersion medium
  538. OPDg = tg * np.sqrt(ng * ng - NA * NA * rho * rho) - tg0 * np.sqrt(ng0 * ng0 - NA * NA * rho * rho) # OPD in the coverslip
  539. W = 2 * np.pi / wavelength * (OPDs + OPDi + OPDg)
  540. phase = np.cos(W) + 1j * np.sin(W)
  541. J = scipy.special.jv(0, scaling_factor.reshape(-1, 1) * rho)
  542. C, residuals, _, _ = np.linalg.lstsq(J.T, phase.T, rcond=-1)
  543. b = 2 * np.pi * r.reshape(-1, 1) * NA / wavelength
  544. J0 = lambda x: scipy.special.j0(x)
  545. J1 = lambda x: scipy.special.j1(x)
  546. denom = scaling_factor * scaling_factor - b * b
  547. R = scaling_factor * J1(scaling_factor * a) * J0(b * a) * a - b * J0(scaling_factor * a) * J1(b * a) * a
  548. R /= denom
  549. PSF_rz = (np.abs(R.dot(C))**2).T
  550. self.PSF_rz = PSF_rz / np.max(PSF_rz)
  552. def __call__(self):
  553. return self.PSF_rz
  556. class Noise(nn.Module):
  557. def __init__(self, params):
  558. super().__init__()
  559. self.sig_eps = params["sig_eps"]
  560. self.a = params["poisson_weight"]
  562. def forward(self, x):
  563. x = x + torch.randn_like(x) * (x * self.a + self.sig_eps)
  564. return x
  567. class PreProcess(nn.Module):
  568. def __init__(self, min, max, params):
  569. super().__init__()
  570. self.min = min
  571. self.max = max
  572. self.background = torch.tensor(params["background"])
  573. self.gamma = dist.Gamma(torch.tensor([1.0]),
  574. torch.tensor(params["background"]))
  576. def forward(self, x):
  577. x = torch.clip(x, min=self.min, max=self.max)
  578. #x = (x - self.min) / (self.max - self.min)
  579. return x
  581. def sample(self, x):
  582. x = (x - self.min) / (self.max - self.min)
  583. # x = x + self.background
  584. # max_value = torch.quantile(x.flatten(), self.max)
  585. x = torch.clip(x, min=self.min, max=self.max) #max_value.item())
  586. return x
  589. class Hill(nn.Module):
  590. def __init__(self, n, ka, params):
  591. super().__init__()
  592. self.n_init = n
  593. self.ka_init = ka
  594. self.n = n
  595. self.k = ka
  596. self.x1 = torch.ones(1).to(device=params["device"]) * 0.5
  597. self.x2 = torch.ones(1).to(device=params["device"]) * 0.8
  599. def forward(self, x):
  600. n = F.sigmoid(self.n)
  601. ka = torch.clip(self.ka, min=0.)
  602. x = self.hill(x, n, ka)
  603. return x
  605. def hill_with_best_value(self, x):
  606. v = self.solve_hill(
  607. self.find_y(x, self.x1),
  608. self.find_y(x, self.x2),
  609. self.hill_ideal(self.x1),
  610. self.hill_ideal(self.x2),)
  611. x = self.hill(x, v['n'], v["ka"])
  612. return x
  614. def sample(self, x):
  615. x = self.hill(x, self.n_init, self.ka_init)
  616. return x
  618. def hill(self, x, n, ka):
  619. return (ka ** n + 1) * x ** n / (ka ** n + x ** n)
  621. def hill_ideal(self, x):
  622. return x ** 0.5 / (1 + x ** 0.5)
  624. def find_y(self, x, x1):
  625. return torch.quantile(x.flatten(), x1)
  627. def solve_hill(self, x1, x2, y1, y2):
  628. a = torch.log(x2) * (torch.log(1 - y1) - torch.log(y1))
  629. b = torch.log(x1) * (torch.log(1 - y2) - torch.log(y2))
  630. c = torch.log(y2) + torch.log(1 - y1)
  631. d = torch.log(y1) + torch.log(1 - y2)
  632. self.k = torch.exp((a - b) / (c - d))
  633. self.n = (torch.log(1 - y1) - torch.log(y1)) \
  634. / (torch.log(self.k) - torch.log(x1))
  635. return {"n" : self.n,
  636. "ka" : self.k }
  638. def inverse_hill(self, x):
  639. x = x.clip(min = 0., max = 1.)
  640. return ((self.k - x + 1) / (self.k * x)) ** (- 1 / self.n)
  643. class ImagingProcess(nn.Module):
  644. def __init__(self, params):
  645. super().__init__()
  646. self.device = params["device"]
  647. self.mu_z = params["mu_z"]
  648. self.sig_z = params["sig_z"]
  649. self.log_ez0 = nn.Parameter(
  650. (torch.tensor(params["mu_z"] + 0.5 \
  651. * params["sig_z"] ** 2)).to(self.device),
  652. requires_grad=True)
  653. self.emission = Emission()
  654. self.blur = Blur(params = params)
  655. self.noise = Noise(params)
  656. self.preprocess = PreProcess(min=0., max=1., params=params)
  657. self.hill = Hill(n=0.5, ka=1., params=params)
  659. def forward(self, x):
  660. out = self.blur(x)
  661. x = out["out"]
  662. x = self.preprocess(x)
  663. # x = self.hill.sample(x)
  664. out = {"out" : x ,
  665. "psf_loss" : out["psf_loss"] }
  666. return out
  669. class JNet(nn.Module):
  670. def __init__(self, params):
  671. super().__init__()
  672. t1 = time.time()
  673. print('initializing JNet model...')
  674. scale_factor = (params["scale"], 1, 1)
  675. hidden_channels_list = params["hidden_channels_list"].copy()
  676. attn_list = params["attn_list"].copy()
  677. hidden_channels = hidden_channels_list.pop(0)
  678. attn_list.pop(0)
  679. self.prev0 = JNetBlock0(in_channels = 1 ,
  680. out_channels = hidden_channels,)
  681. self.prev = nn.ModuleList(
  682. [JNetBlock(
  683. in_channels = hidden_channels,
  684. hidden_channels = hidden_channels,
  685. dropout = params["dropout"],
  686. ) for _ in range(params["nblocks"])
  687. ])
  688. self.mid = JNetLayer(
  689. in_channels = hidden_channels ,
  690. hidden_channels_list = hidden_channels_list ,
  691. attn_list = attn_list ,
  692. nblocks = params["nblocks"] ,
  693. dropout = params["dropout"] ,
  694. ) if hidden_channels_list else nn.Identity()
  695. self.postx = nn.ModuleList(
  696. [JNetBlock(in_channels = hidden_channels,
  697. hidden_channels = hidden_channels,
  698. dropout = params["dropout"],
  699. ) for _ in range(params["nblocks"])
  700. ])
  701. self.postx.append(JNetBlockN(
  702. in_channels = hidden_channels ,
  703. out_channels = 1 ,
  704. ))
  705. self.postz = nn.ModuleList(
  706. [JNetBlock(in_channels = hidden_channels,
  707. hidden_channels = hidden_channels,
  708. dropout = params["dropout"],
  709. ) for _ in range(params["nblocks"])
  710. ])
  711. self.postz.append(JNetBlockN(
  712. in_channels = hidden_channels ,
  713. out_channels = 1 ,
  714. ))
  715. self.image = ImagingProcess(params = params)
  716. self.upsample = JNetUpsample(scale_factor = scale_factor)
  717. self.activation = params["activation"]
  718. self.superres = params["superres"]
  719. self.reconstruct = params["reconstruct"]
  720. self.apply_vq = params["apply_vq"]
  721. self.vq = VectorQuantizer(threshold=params["threshold"],
  722. device=params["device"])
  723. t2 = time.time()
  724. print(f'JNet init done ({t2-t1:.2f} s)')
  725. self.use_x_quantized = params["use_x_quantized"]
  726. self.tau = 1.
  727. a = params["poisson_weight"]
  728. a_inv_sig = np.log(a / (1 - a))
  729. self.a = nn.Parameter(torch.tensor(a_inv_sig)
  730. .to(device=params["device"]))
  731. s = params["sig_eps"]
  732. s_inv_sig = np.log(s / (1 - s))
  733. self.sigma = nn.Parameter(torch.tensor(s_inv_sig)
  734. .to(device=params["device"]))
  736. def forward(self, x):
  737. if self.superres:
  738. x = self.upsample(x)
  739. x = self.prev0(x)
  740. for f in self.prev:
  741. x = f(x)
  742. _x = self.mid(x)
  743. x = _x
  744. z = _x
  745. for f in self.postx:
  746. x = f(x)
  747. for f in self.postz:
  748. z = f(z)
  749. if self.apply_vq:
  750. if self.use_x_quantized:
  751. x, qloss = self.vq(x)
  752. else:
  753. _, qloss = self.vq(x)
  754. if self.reconstruct:
  755. lu = x * z
  756. out = self.image(lu)
  757. r = out["out"]
  758. psf_loss = out["psf_loss"]
  759. else:
  760. r = x
  762. out = {"enhanced_image" : x ,
  763. "reconstruction" : r ,
  764. "mid" : _x ,
  765. "estim_luminance" : z ,
  766. "poisson_weight" : F.sigmoid(self.a) ,
  767. "gaussian_sigma" : F.sigmoid(self.sigma),
  768. }
  769. vqd = {"quantized_loss" : qloss} if self.apply_vq\
  770. else {"quantized_loss" : None}
  771. out = dict(**out, **vqd)
  772. psl = {"psf_loss": psf_loss} if self.reconstruct\
  773. else {"psf_loss": None}
  774. out = dict(**out, **psl)
  775. return out

model.py at commit 4d6bc0e, no license · at the source

Overview

Authors: Haruhiko Morita1, Shuto Hayashi1,2, Takahiro Tsuji3, Daisuke Kato4, Hiroaki Wake5,6, Teppei Shimamura1
ORCID iDs: Haruhiko Morita
  1. Department of Computational and Systems Biology, Division of Biological Data Science, Medical Research Laboratory, Institute of Integrated Research, Institute of Science Tokyo, Yushima, Bunkyoku, Tokyo 113-8510, Japan
  2. Innovation Center of NanoMedicine, Kawasaki Institute of Industrial Promotion, 3-25-14 Tonomachi, Kawasaki-ku, Kawasaki 210-0821, Japan
  3. Nagoya University Graduate School of Medicine, Department of Anatomy and Molecular Cell Biology, Tsurumaicho, Showa-ku, Nagoya, Aichi 466-8560, Japan
  4. Department of Physiology, Nippon Medical School Graduate School of Medicine, 1-25-16 Nezu, Bunkyo-ku, Tokyo 113-0031, Japan
  5. Department of Anatomy and Molecular Cell Biology, Nagoya University Graduate School of Medicine, 65 Tsurumai-cho, Showa-ku, Nagoya 466-8550, Japan
  6. Division of Multicellular Circuit Dynamics, National Institute for Physiological Sciences, National Institute of Natural Sciences, 38 Nishigonaka, Myodaiji-cho, Okazaki, Japan
Institutions: Institute of Science Tokyo (Japan); Innovation Center of NanoMedicine (Japan); Nagoya University (Japan); Nippon Medical School (Japan); National Institute for Physiological Sciences (Japan)
Journal: Cell reports methods, volume 6, issue 8, article 101476
Dates: received 30 September 2025; accepted 6 May 2026; published online 5 June 2026; in print August 2026
Type: Research article · Language: English
License: CC BY-NC
Identifiers: DOI 10.1016/j.crmeth.2026.101476 · PMID 42248144 · PMCID PMC13494536 · OpenAlex W4415252322
Open access: gold, a free copy (OpenAlex)
Preprint: doi:10.2139/ssrn.5561404
Status: code verified
Categories: histology / microscopy (modality), optical imaging (calcium, voltage, 2-photon) (modality), human (organism), mouse (organism), other condition (population), cellular / molecular (subfield)
Methods: Connectivity, Spectral & time-frequency, Evoked potentials, Machine learning, fMRI & imaging
Keywords: multiphoton excitation microscopy, unsupervised machine learning, autoencoder, microglia
MeSH: Deep Learning*, Microscopy, Fluorescence, Multiphoton*, Unsupervised Machine Learning*, Algorithms, Animals, Brain Neoplasms, Humans, Image Processing, Computer-Assisted, Imaging, Three-Dimensional, Mice, Microglia (* major topic)
Topic: Cell Image Analysis Techniques (Biophysics, Biochemistry, Genetics and Molecular Biology), according to OpenAlex
Funding: Japan Agency for Medical Research and Development (JP26zf0127012, JP26tm0424226); Japan Society for the Promotion of Science (22H04925, 26K03026, 23H04938); Core Research for Evolutional Science and Technology (JP26gm2010002, JP26wm0625519, JP25wm0325068, JP26nk0101112); Japan Science and Technology Agency Moonshot Research and Development Program (JPMJMS2026); Japan Science and Technology Agency; University of Tokyo
Citations: not cited yet (Europe PMC); 36 references in the paper
Research resources: RRID:IMSR_JAX:005582

Abstract

Two-photon microscopy enables the non-invasive imaging of deep living tissue. Quantitative three-dimensional analysis is hampered by axial blur and anisotropic resolution in two-photon microscopy. We introduce the two-photon microscopy image enhancement network (TENET), a fully unsupervised framework that simultaneously performs deblurring, up to 12× resolution enhancement, and semantic segmentation on two-photon microscopy volumes in a single pass. TENET embeds a physics-informed blur-generation module with a trainable neural implicit point spread function (PSF), requiring only approximate PSF initialization rather than rigorous experimental measurement, paired “unblurred” images, or isotropy assumptions. On synthetic images, fluorescent beads, and in vivo microglia datasets, TENET surpasses the total-variation-regularized Richardson-Lucy (RLTV) algorithm, CARE, and Neuroclear in image fidelity and segmentation accuracy. Using time-lapse images of microglia surrounding metastatic brain tumors, TENET enables automated 3D morphometry that reveals dynamic microglia-tumor interactions. By converting blur-limited two-photon microscopy data into high-fidelity volumetric reconstructions with ready-to-use masks, TENET streamlines downstream analysis and expands the reach of deep-tissue live imaging.

Reproduced under the paper's license (CC BY-NC), from the paper cited above.

Repositories

Its files are read in the Code ↔ Paper reader above, with 5 matches between paragraphs and lines of code.

HaLU-9000/TENET

License: none: the authors keep all their rights
State: the link answers, verified on 27 September 2026
Evidence: files inventoried
Commit: 4d6bc0e988969e3ca01c3f4fece7f8d96d6d9867, 27 May 2026
Languages: Python (15), Jupyter (6), Shell (2)
Size: 41 files, 23 scripts
Software Heritage: not archived
Found in: the resources table
Holds: README, environment (requirements.txt), 6 notebooks
Not found: license file, CITATION.cff, tests, continuous integration, documentation
Tools: PyTorch (19 files), NumPy (18 files), Matplotlib (11 files), tifffile (10 files), OpenCV (5 files), pandas (4 files), scikit-image (3 files), SciPy (3 files)
Availability: 1 check, the latest on 27 September 2026: the link answers
  • 27 September 2026: the link answers
24 files
At the source: github.com/HaLU-9000/TENET

Zenodo 19663648

License: CC-BY-4.0
State: the link answers, verified on 27 September 2026
Evidence: files inventoried
Size: 1 file
Software Heritage: not checked
Found in: the resources table
Not found: README, license file, CITATION.cff, environment file, tests, continuous integration, documentation
Tools: PyTorch (18 files), NumPy (17 files), Matplotlib (11 files), tifffile (9 files), OpenCV (5 files), pandas (4 files), scikit-image (3 files), SciPy (3 files)
Availability: 1 check, the latest on 27 September 2026: the link answers (HTTP 200)
  • 27 September 2026: the link answers (HTTP 200)
23 files
At the source: doi.org/10.5281/zenodo.19663648

The paper's code and data availability statement is in the Data section.

Tracing map

Proposed by the machine: these links were found in the paper and verified at the source, without human review. The map will receive a Zenodo DOI once one of the paper's authors has validated it with their ORCID.

What the map holds:

  • 2 repositories of the authors' code, each at its verified commit, with its license and how the link was found in the paper;
  • 45 scripts, each with its path and the digest of its content;
  • 5 matches between paragraphs of the paper and lines of the code (method lexical-v1);
  • neither the text of the paper nor the code itself.

Its JSON (tracing-map.json) is deposited on Zenodo with its DOI once the map is validated.

Data

Datasets cited

Data and code availability

The code for reproducing the tables and figures is available at https://github.com/HaLU-9000/TENET/. The data used for training and raw data of the figures are available at https://zenodo.org/records/18869333 and https://zenodo.org/records/15545449. These repositories include all scripts, model weights, and representative datasets necessary to reproduce the results presented in this manuscript.

Reproduced under the paper's license (CC BY-NC), from the paper cited above.

Versions

The history of this record: each version stored by the harvester or made by a correction of its authors or of the maintainers of its code, and what changed in its facts. The texts of the paper (its abstract, its availability statements) are not part of it; versions that changed only those are not listed.

Version 2, 28 September 2026

  • Authors: added Haruhiko Morita (0009-0004-4639-6175); removed Haruhiko Morita

Version 1, 27 September 2026: the first record

Recorded: type, language, journal, volume, issue, pages, dates, 6 authors, 4 keywords, 11 MeSH terms, 6 funders, 36 references, 1 RRID.

Cite

This paper

Morita, H., Hayashi, S., Tsuji, T., Kato, D., Wake, H., & Shimamura, T. (2026). Unsupervised deep learning enables blur-free resolution enhancement in two-photon microscopy. Cell reports methods, 6(8), 101476. https://doi.org/10.1016/j.crmeth.2026.101476

BibTeX

@article{morita2026unsupervised,
author = {Morita, Haruhiko and Hayashi, Shuto and Tsuji, Takahiro and Kato, Daisuke and Wake, Hiroaki and Shimamura, Teppei},
title = {{Unsupervised deep learning enables blur-free resolution enhancement in two-photon microscopy}},
journal = {Cell reports methods},
year = {2026},
month = jun,
volume = {6},
number = {8},
pages = {101476},
publisher = {Elsevier},
issn = {2667-2375},
doi = {10.1016/j.crmeth.2026.101476},
url = {https://doi.org/10.1016/j.crmeth.2026.101476},
pmid = {42248144},
pmcid = {PMC13494536}
}

RIS

TY - JOUR
AU - Morita, Haruhiko
AU - Hayashi, Shuto
AU - Tsuji, Takahiro
AU - Kato, Daisuke
AU - Wake, Hiroaki
AU - Shimamura, Teppei
TI - Unsupervised deep learning enables blur-free resolution enhancement in two-photon microscopy
T2 - Cell reports methods
J2 - Cell Rep Methods
PY - 2026
DA - 2026/06/05
VL - 6
IS - 8
SP - 101476
SN - 2667-2375
PB - Elsevier
DO - 10.1016/j.crmeth.2026.101476
UR - https://doi.org/10.1016/j.crmeth.2026.101476
LA - en
ER -

CSL-JSON

{
"id": "10.1016/j.crmeth.2026.101476",
"type": "article-journal",
"title": "Unsupervised deep learning enables blur-free resolution enhancement in two-photon microscopy",
"container-title": "Cell reports methods",
"author": [
{
"family": "Morita",
"given": "Haruhiko"
},
{
"family": "Hayashi",
"given": "Shuto"
},
{
"family": "Tsuji",
"given": "Takahiro"
},
{
"family": "Kato",
"given": "Daisuke"
},
{
"family": "Wake",
"given": "Hiroaki"
},
{
"family": "Shimamura",
"given": "Teppei"
}
],
"container-title-short": "Cell Rep Methods",
"volume": "6",
"issue": "8",
"page": "101476",
"DOI": "10.1016/j.crmeth.2026.101476",
"PMID": "42248144",
"PMCID": "PMC13494536",
"ISSN": "2667-2375",
"publisher": "Elsevier",
"URL": "https://doi.org/10.1016/j.crmeth.2026.101476",
"language": "en",
"issued": {
"date-parts": [
[
2026,
6,
5
]
]
}
}

The tracing map gets a citation of its own once an author has validated it and it has a DOI.

Similar papers

The papers with a page that share the most with this one: the tools found in their code, their categories, datasets, cited references and authors, the rarest counting most.

[1] doi:10.1016/j.isci.2026.117010 [code]
Deep learning-assisted mapping of dendritic spines using sequential 2D two-photon calcium imaging.
Journal: iScience
In common: tifffile, OpenCV, scikit-image, 5 other tools, optical imaging (calcium, voltage, 2-photon), histology / microscopy, cellular / molecular, 1 reference
[2] doi:10.7554/elife.105081 [code]
Movie reconstruction from mouse visual cortex activity.
Journal: eLife
In common: tifffile, OpenCV, PyTorch, 4 other tools, optical imaging (calcium, voltage, 2-photon), mouse, 2 references
[3] doi:10.1038/s41598-026-57519-w [code]
Automated segmentation of neurons and spinal cord structures in immunofluorescence images using SpineDL.
Journal: Scientific reports
In common: tifffile, OpenCV, scikit-image, 5 other tools, histology / microscopy, other condition, mouse
[4] doi:10.1038/s41467-026-73389-2 [code]
Physics-informed multi-encoder adaptive optics enables rapid aberration correction for intravital microscopy of deep complex tissue.
Journal: Nature communications
In common: OpenCV, SciPy, Matplotlib, 1 other tool, optical imaging (calcium, voltage, 2-photon), histology / microscopy, mouse, 3 references
[5] doi:10.1038/s41467-026-72709-w [code]
An epifluorescence microscope design for naturalistic behavior and cellular activity in freely moving Caenorhabditis elegans.
Journal: Nature communications
In common: tifffile, OpenCV, scikit-image, 5 other tools, histology / microscopy, cellular / molecular
[6] doi:10.1038/s41467-026-73045-9 [code]
Aberration-aware 3D localization microscopy via self-supervised neural-physics learning.
Journal: Nature communications
In common: tifffile, OpenCV, scikit-image, 5 other tools, histology / microscopy, cellular / molecular
[7] doi:10.1364/boe.600665 [code]
NeuroSeg-MF: robust neuron segmentation in two-photon Ca&lt;sup&gt;2+&lt;/sup&gt; imaging using multi-feature fusion and detection-guided SAM.
Journal: Biomedical optics express
In common: tifffile, OpenCV, PyTorch, 4 other tools, optical imaging (calcium, voltage, 2-photon), 1 reference
[8] doi:10.1371/journal.pcbi.1013441 [code]
Large vision model framework for automated C. elegans analysis: From static morphometry to dynamic neural activity.
Journal: PLoS computational biology
In common: tifffile, OpenCV, scikit-image, 5 other tools, optical imaging (calcium, voltage, 2-photon)
[9] doi:10.1016/j.isci.2026.116206 [code]
Gut distension evokes rapid neural dynamics in vagal and hindbrain populations of larval zebrafish.
Journal: iScience
In common: tifffile, OpenCV, scikit-image, 5 other tools, optical imaging (calcium, voltage, 2-photon)
[10] doi:10.1002/jbio.70260 [code]
Cell-MICS: Detecting Immune Cells With Label-Free Two-Photon Autofluorescence and Deep Learning.
Journal: Journal of biophotonics
In common: tifffile, PyTorch, pandas, 3 other tools, optical imaging (calcium, voltage, 2-photon), histology / microscopy, cellular / molecular, 1 reference

Contribute

The authors of this paper can claim it, correct its record and validate its tracing map, and the maintainers of its code (its owner, or a public member of its organization) correct what it says of their repository; anyone signed in can ask for its removal. Every request goes to OSCR's own machine, which answers it; your account page follows them.

Sign in with ORCID to claim this paper as one of its authors, correct its record or validate its tracing map: when the paper's metadata lists your ORCID iD, you are recognized at once. Maintainers of its code: sign in with GitHub, then claim the repository on your account page.

Request its removal

To ask OSCR to remove this record, the copies of its authors' scripts or its tracing map, use the removal request page: signed in, you say who you are, what to remove and why, then review and confirm the request. Published rules decide every request (how).

Discussion, reproductions, activity

Discussion: questions and error reports about this paper and its code, from signed-in readers and its authors. It opens with sign-in.

Reproductions: reports from readers who ran the authors' code: what they reproduced, with which environment, commit and data. It opens with sign-in.

Activity: what happens around this paper: new versions of its record, its map's validation, discussions and reproductions. It opens with sign-in.