radon.py

import torch
from torch import nn
import torch.nn.functional as F

from pytorch_radon.utils import PI, SQRT2, deg2rad, affine_grid, grid_sample
from pytorch_radon.filters import RampFilter

class Radon(nn.Module):
    def __init__(self, in_size=None, theta=None, circle=True, dtype=torch.float):
        super(Radon, self).__init__()
        self.circle = circle
        self.theta = theta
        if theta is None:
            self.theta = torch.arange(180)
        self.dtype = dtype
        self.all_grids = None
        if in_size is not None:
            self.all_grids = self._create_grids(self.theta, in_size, circle)

    def forward(self, x):
        N, C, W, H = x.shape
        assert W == H

        if self.all_grids is None:
            self.all_grids = self._create_grids(self.theta, W, self.circle)

        if not self.circle:
            diagonal = SQRT2 * W
            pad = int((diagonal - W).ceil())
            new_center = (W + pad) // 2
            old_center = W // 2
            pad_before = new_center - old_center
            pad_width = (pad_before, pad - pad_before)
            x = F.pad(x, (pad_width[0], pad_width[1], pad_width[0], pad_width[1]))

        N, C, W, _ = x.shape
        out = torch.zeros(N, C, W, len(self.theta), device=x.device, dtype=self.dtype)

        for i in range(len(self.theta)):
            rotated = grid_sample(x, self.all_grids[i].repeat(N, 1, 1, 1).to(x.device))
            out[..., i] = rotated.sum(2)

        return out

    def _create_grids(self, angles, grid_size, circle):
        if not circle:
            grid_size = int((SQRT2*grid_size).ceil())
        grid_shape = [1, 1, grid_size, grid_size]
        all_grids = []
        for theta in angles:
            theta = deg2rad(theta, self.dtype)
            R = torch.tensor([[
                [theta.cos(), theta.sin(), 0],
                [-theta.sin(), theta.cos(), 0],
            ]], dtype=self.dtype)
            all_grids.append(affine_grid(R, grid_shape))
        return all_grids
class IRadon(nn.Module):
    def __init__(self, in_size=None, theta=None, circle=True,
                 use_filter=RampFilter(), out_size=None, dtype=torch.float):
        super(IRadon, self).__init__()
        self.circle = circle
        self.theta = theta if theta is not None else torch.arange(180)
        self.out_size = out_size
        self.in_size = in_size
        self.dtype = dtype
        self.deg2rad = lambda x: deg2rad(x, dtype)
        self.ygrid, self.xgrid, self.all_grids = None, None, None
        if in_size is not None:
            self.ygrid, self.xgrid = self._create_yxgrid(in_size, circle)
            self.all_grids = self._create_grids(self.theta, in_size, circle)
        self.filter = use_filter if use_filter is not None else lambda x: x

    def forward(self, x):
        it_size = x.shape[2]
        ch_size = x.shape[1]

        if self.in_size is None:
            self.in_size = int((it_size/SQRT2).floor()) if not self.circle else it_size
        if None in [self.ygrid, self.xgrid, self.all_grids]:
            self.ygrid, self.xgrid = self._create_yxgrid(self.in_size, self.circle)
            self.all_grids = self._create_grids(self.theta, self.in_size, self.circle)
        print('x.shape', x.shape)
        x = self.filter(x)

        reco = torch.zeros(x.shape[0], ch_size, it_size, it_size, device=x.device, dtype=self.dtype)

        print('')
        print('x.shape', x.shape)

        for i_theta in range(len(self.theta)):
            print(f'self.all_grids[{i_theta}].shape', self.all_grids[i_theta].shape)
            reco += grid_sample(x, self.all_grids[i_theta].to(x.device))

        if not self.circle:
            W = self.in_size
            diagonal = it_size
            pad = int(torch.tensor(diagonal - W, dtype=torch.float).ceil())
            new_center = (W + pad) // 2
            old_center = W // 2
            pad_before = new_center - old_center
            pad_width = (pad_before, pad - pad_before)
            reco = F.pad(reco, (-pad_width[0], -pad_width[1], -pad_width[0], -pad_width[1]))

        if self.circle:
            reconstruction_circle = (self.xgrid ** 2 + self.ygrid ** 2) <= 1
            reconstruction_circle = reconstruction_circle.repeat(x.shape[0], ch_size, 1, 1)
            reco[~reconstruction_circle] = 0.

        reco *= PI.to(reco.device)/(2*len(self.theta))

        if self.out_size is not None:
            pad = (self.out_size - self.in_size)//2
            reco = F.pad(reco, (pad, pad, pad, pad))

        return reco

    def _create_yxgrid(self, in_size, circle):
        if not circle:
            in_size = int((SQRT2*in_size).ceil())
        unitrange = torch.linspace(-1, 1, in_size, dtype=self.dtype)
        return torch.meshgrid(unitrange, unitrange)

    def _xy_to_t(self, theta):
        return self.xgrid*self.deg2rad(theta).cos() - self.ygrid*self.deg2rad(theta).sin()

    def _create_grids(self, angles, grid_size, circle):
        if not circle:
            grid_size = int((SQRT2*grid_size).ceil())
        all_grids = []
        for i_theta, theta in enumerate(angles):
            X = torch.ones([grid_size]*2, dtype=self.dtype)*i_theta*2./(len(angles)-1)-1.
            Y = self._xy_to_t(theta)
            all_grids.append(torch.stack((X, Y), dim=-1).unsqueeze(0))
        return all_grids

import skimage.data as d
import numpy as np
import matplotlib.pyplot as plt
import torch as th

a = d.camera().astype(np.float32)

fig, ax = plt.subplots()
ax.imshow(a)
plt.show()

img = th.as_tensor(a).unsqueeze(0).unsqueeze(0)
circle = False
theta = torch.arange(360)
r = Radon(img.shape[2], theta, circle)
ir = IRadon(img.shape[2], theta, circle)
sino = r(img)
reco = ir(sino)
fig, ax = plt.subplots()
ax.imshow(sino.squeeze())
plt.show()