In [ ]:
%matplotlib inline

In [ ]:
# coding: utf-8

# TODO: replace paths to mmaped files

import time
import sys
from skimage import io
import glob
import astra
import numpy as np
import pylab as plt
import numexpr
from pprint import pprint

In [ ]:
plt.viridis()

In [ ]:
def create_test_sinograms():
    '''
    Create a simple hollow cube phantom

    :param size: Cube size length (int >0)
    :return: numpy mmaped array cube
    '''
    
    files = sorted(glob.glob(
        '/diskmnt/a/makov/yaivan/tomography_dbg/projections/MMC1_2.82um__projection*.tiff'))
#         '/diskmnt/a/makov/yaivan/tomography_dbg1/projections/MMC1_2.82um__projection*.tiff'))
    print(len(files))
    f0 = io.imread(files[0])
#     print(f0.shape)
#     plt.figure()
#     plt.imshow(f0, interpolation='nearest')
#     plt.axis('tight')
#     cube = np.zeros((size,size,size), dtype='float32')
    cube = np.memmap('/diskmnt/fast/makov/sinogram.tmp',
                     dtype='float32', mode='w+', shape=(f0.shape[0], len(files), f0.shape[1]))
    for ifd, f in enumerate(files):
        t=io.imread(f)
        cube[:,ifd,:] = t #[20:,...]

#     plt.figure()
#     plt.imshow(cube[2], interpolation='nearest')
#     plt.axis('tight')
    return cube

# create_test_sinograms()

In [ ]:
def create_test_cube(size):
    '''
    Create a simple hollow cube phantom

    :param size: Cube size length (int >0)
    :return: numpy mmaped array cube
    '''
#     cube = np.zeros((size,size,size), dtype='float32')
    cube = np.memmap('/diskmnt/fast/makov/cube.tmp',
                     dtype='float32', mode='w+', shape=(size, size, size))
    x0 = int(128.*size/1024)
    x1 = int(895.*size/1024)
    y0 = int(256.*size/1024)
    y1 = int(767.*size/1024)
    cube[x0:x1, x0:x1, x0:x1] = 1
    cube[y0:y1, y0:y1, y0:y1, ] = 0
    return cube




def build_geometry(angles, rec_size):
    """
    Build astra circular geometry for square detector.  See example #5 from python samples

    :params angles_count:List of angles in radians
    :param rec_size: Size of reconstruction volume
    """

    # All distances in [pixels] (MMC1 dataset configuration)

    pixel_size = (2.82473e-3)*8
    os_distance = 56.135 / pixel_size  # object-sample distance
    ds_distance = 225.082 / pixel_size  # detector-sample distance
#     detector_size = rec_size

    # Circular

    # Parameters: width of detector column, height of detector row, #rows, #columns,
    #             angles, distance source-origin, distance origin-detector

    proj_geom = astra.create_proj_geom('cone', ds_distance / os_distance, ds_distance / os_distance,
                                       rec_size[0], rec_size[1], angles,
                                       os_distance, (ds_distance - os_distance))
    return proj_geom


def create_sino3d_gpu_mm(data, proj_geom, vol_geom, returnData=True, gpuIndex=None, sinogram_volume=None):
    """Create a forward projection of an image (3D).

    :param data: Image data or ID.
    :type data: :class:`numpy.ndarray` or :class:`int`
    :param proj_geom: Projection geometry.
    :type proj_geom: :class:`dict`
    :param vol_geom: Volume geometry.
    :type vol_geom: :class:`dict`
    :param returnData: If False, only return the ID of the forward projection.
    :type returnData: :class:`bool`
    :param gpuIndex: Optional GPU index.
    :type gpuIndex: :class:`int`
    :returns: :class:`int` or (:class:`int`, :class:`numpy.ndarray`) -- If ``returnData=False``, returns the ID of the forward projection. Otherwise, returns a tuple containing the ID of the forward projection and the forward projection itself, in that order.

    """

    if isinstance(data, np.ndarray):
        volume_id = astra.data3d.link('-vol', vol_geom, data)
    else:
        volume_id = data

    if isinstance(sinogram_volume, np.ndarray):
        sino_id = astra.data3d.link('-sino', proj_geom, sinogram_volume)
    else:
        sino_id = astra.data3d.create('-sino', proj_geom, 0)

    algString = 'FP3D_CUDA'
    cfg = astra.astra_dict(algString)
    if gpuIndex is not None:
        cfg['option'] = {'GPUindex': gpuIndex}
    cfg['ProjectionDataId'] = sino_id
    cfg['VolumeDataId'] = volume_id
    alg_id = astra.algorithm.create(cfg)
    astra.algorithm.run(alg_id)
    astra.algorithm.delete(alg_id)

    if isinstance(data, np.ndarray):
        astra.data3d.delete(volume_id)
    if returnData:
        return sino_id, astra.data3d.get(sino_id)
    else:
        return sino_id


def test_reconstruction_cone_fdk(gpus_list):
    """
    Build test object, sinogram and reconstrucion
    See http://www.astra-toolbox.com/docs/index.html for other odcs

    :param size: Size of reconstruction volume
    :params angles_count: Number of angles in range 0-PI
    :params gpus_list: list of GPU to use for reconstruction, usually [0,] or [1,] or [0,1]
    """
    # Set up multi-GPU usage.
    # This only works for 3D GPU forward projection and back projection.
    astra.astra.set_gpu_index(gpus_list)


#     angles = angles-(angles.max()+angles.min())/2
    print('Creating sinogram')
    t = time.time()
    sinogram_volume = create_test_sinograms()
    print(sinogram_volume.shape)
#     angles = np.linspace(0, np.pi, angles_count, False)
    angles = np.arange(sinogram_volume.shape[1])*0.1*4/180*np.pi
    proj_geom = build_geometry(angles, (sinogram_volume.shape[0],sinogram_volume.shape[2]))
    proj_id = astra.data3d.link('-sino', proj_geom, sinogram_volume)
    print(time.time() - t)
    

    print('Creating phantom')
    t = time.time()
    # Create a simple hollow cube phantom
#     cube = create_test_cube(size)

    # Create projection geometry
    

    # Create volume desscription
    vol_geom = astra.create_vol_geom((sinogram_volume.shape[2],sinogram_volume.shape[2],sinogram_volume.shape[0]))
#     vol_geom['option']['WindowMinZ']=  -0.5
#     vol_geom['option']['WindowMaxZ']=  0.5
    pprint(vol_geom)
#     # Upload data to ASTRA
#     cube_id = astra.data3d.link('-vol', vol_geom, cube)

    print(time.time() - t)

    # Create mmapfile for sinogram
#     sinogram_volume = np.memmap('/diskmnt/fast/makov/sinogram.tmp',
#                                 dtype='float32', mode='w+',
#                                 shape=(cube.shape[0], len(angles), cube.shape[1]))  # TODO: check for non cube
    


    # Create sinogram from cube and NOT send data to scipt
#     proj_id = create_sino3d_gpu_mm(cube_id, proj_geom, vol_geom,
#                                    returnData=False, sinogram_volume=sinogram_volume)

#     astra.data3d.delete(cube_id)
#     del cube


    print('Creating volume')
    t = time.time()

    # To get data by id
    # projection_data = astra.data3d.get(proj_id)
    rec_volume = np.memmap('/diskmnt/fast/makov/rec.tmp',
                           dtype='float32', mode='w+',
                           shape=(vol_geom['GridSliceCount'],  # TODO: check for non cube
                                  vol_geom['GridRowCount'],   
                                  vol_geom['GridColCount'])
#                             shape=(5,500,500)
                           )
    print(rec_volume.shape)
    rec_id = astra.data3d.link('-vol', vol_geom, rec_volume)
    # Set up the parameters for a reconstruction algorithm using the GPU
    # Complete list of suporteed algoritms can be found
    # http://www.astra-toolbox.com/docs/algs/index.html
    print(time.time() - t)
    print('Start reconstruction')
    t = time.time()
    cfg = astra.astra_dict('FDK_CUDA')
#     cfg = astra.astra_dict('FDK_CUDA')
    cfg['ReconstructionDataId'] = rec_id
    cfg['ProjectionDataId'] = proj_id
    cfg['option'] = {'ShortScan': True}

    # Create the algorithm object from the configuration structure
    alg_id = astra.algorithm.create(cfg)

    # For FDK number of iterationa my be 1
    astra.algorithm.run(alg_id,1000)

    print('Stop reconstruction')
    print(time.time() - t)
#     rec = astra.data3d.get(rec_id)
#     plt.figure()
#     plt.imshow(rec[:,:,int(rec.shape[-1]/2)], cmap=plt.cm.jet)
#     plt.show()

    # Clean up. Note that GPU memory is tied up in the algorithm object,
    # and main RAM in the data objects.
    astra.data3d.info()
    astra.algorithm.delete(alg_id)
    astra.data3d.delete(proj_id)
    astra.data3d.delete(rec_id)
#     del rec_volume
#     del sinogram_volume
    return rec_volume,sinogram_volume, angles, proj_geom

In [ ]:
t = time.time()
rec_volume,sinogram_volume, angles, proj_geom = test_reconstruction_cone_fdk(gpus_list=[0,1])
print(time.time() - t)

# # load and show reconstruction cut
# rec_volume = np.memmap('/diskmnt/fast/makov/rec.tmp',
#                        dtype='float32', mode='r',
#                        shape=(200, 200, 200)
#                        )

plt.figure(figsize=(15,10))
plt.imshow(rec_volume[2],
           cmap=plt.cm.gray)

plt.colorbar()
plt.show()

In [ ]:
vol_geom

In [ ]:
test_rec = io.imread('/diskmnt/a/makov/yaivan/tomography_dbg/slices/MMC1_2.82um__rec0120.tiff')
# test_sino = io.imread('/diskmnt/a/makov/yaivan/tomography_dbg1/sinograms/MMC1_2.82um__sino0120.tiff')

In [ ]:
rec_volume.shape

In [ ]:
plt.figure(figsize=(15,10))
plt.imshow(rec_volume[2]-test_rec,
           cmap=plt.cm.gray)

plt.colorbar()
plt.show()

In [ ]:
plt.figure(figsize=(15,10))
plt.imshow(test_rec,
           cmap=plt.cm.gray)

plt.colorbar()
plt.show()

In [ ]:
print np.sum(test_rec)
print np.sum(rec_volume[2])

In [ ]:
print np.sum(rec_volume[2])

In [ ]:
print np.sum(rec_volume[2])

In [ ]:
plt.plot(test_sino.sum(axis=1))
plt.grid()
print(np.mean(test_sino.sum(axis=1)))
print(np.median(test_sino.sum(axis=1)))

In [ ]:
def performFP(proj_id, rec_id):
    cfg = astra.astra_dict('FP3D_CUDA')
    cfg['ProjectionDataId'] = proj_id
    cfg['VolumeDataId']     = rec_id
    alg_id = astra.algorithm.create(cfg)
    astra.algorithm.run(alg_id)
    astra.algorithm.delete(alg_id)

def performBP(proj_id, rec_id):
    cfg = astra.astra_dict('BP3D_CUDA')
    cfg['ProjectionDataId']     = proj_id
    cfg['ReconstructionDataId'] = rec_id
    alg_id = astra.algorithm.create(cfg)
    astra.algorithm.run(alg_id)
    astra.algorithm.delete(alg_id)

In [ ]:
print(rec_volume.shape)
slices_num = rec_volume.shape[0]
center = slices_num//2
print slices_num, center
plt.figure(figsize=(15,10))
for sn in [center-1]:
    test_volume = np.require(np.swapaxes(rec_volume[sn:-sn],0,2),requirements=['A','C'])
    test_volume = rec_volume[sn:-sn]
    print(test_volume.shape)
    vol_geom = astra.create_vol_geom((test_volume.shape[2],test_volume.shape[2],test_volume.shape[0]))
    pprint(vol_geom)
    angles = np.arange(sinogram_volume.shape[1])*0.1*4/180*np.pi
    proj_geom = build_geometry(angles, (test_volume.shape[0],test_volume.shape[2]))
    
    rec_id = astra.data3d.link('-vol', vol_geom, test_volume)
    
    proj_id = astra.data3d.create('-sino', proj_geom)
    performFP(proj_id, rec_id)
    sino = astra.data3d.get(proj_id)
    astra.data3d.clear()
#     plt.plot(sino[sino.shape[0]//2][100], label = 'Slices: {}'.format(center-sn))
    plt.plot(sino[sino.shape[0]//2].sum(axis=1), 
             label = 'Calculated sinograms on slices numb: {}'.format(1+(center-sn)))
plt.plot(sinogram_volume[sinogram_volume.shape[0]//2].sum(axis=1), label = 'Original sinogram')
plt.ylim(ymin=0)
plt.grid()
plt.legend(loc=0)
plt.show()

In [ ]:
plt.imshow(sinogram_volume[1:-1][:,1,:], interpolation='nearest')
plt.colorbar(orientation='horizontal')
plt.axis('tight')
plt.title('Original projection')

plt.imshow(sino[:,1,:], interpolation='nearest')
plt.colorbar(orientation='horizontal')
plt.axis('tight')
plt.title('After FDK projection')

In [ ]:
plt.imshow(sinogram_volume[2]-sino[1])
plt.colorbar()

In [ ]:
np.sum(sinogram_volume[2])/np.sum(sino[1])

In [ ]:
sinogram_volume.shape

In [ ]:
cs = sino[:,sino.shape[1]//2,:]

In [ ]:
plt.imshow(sino[sino.shape[0]//2])

In [ ]:
test_volume.shape

In [ ]:
plt.figure(figsize=(15,15))
plt.imshow(sinogram_volume[2]-test_sino,
           cmap=plt.cm.viridis)

plt.colorbar()
plt.show()

In [ ]:
plt.figure(figsize=(15,15))
plt.imshow(test_sino,
           cmap=plt.cm.viridis)

plt.colorbar()
plt.show()

In [ ]:
astra.data3d.info()

In [ ]:
print rec_volume.nbytes/1e9
print sinogram_volume.nbytes/1e9

In [ ]:
def performFP(proj_id, rec_id):
    cfg = astra.astra_dict('FP3D_CUDA')
    cfg['ProjectionDataId'] = proj_id
    cfg['VolumeDataId']     = rec_id
    alg_id = astra.algorithm.create(cfg)
    astra.algorithm.run(alg_id)
    astra.algorithm.delete(alg_id)

def performBP(proj_id, rec_id):
    cfg = astra.astra_dict('BP3D_CUDA')
    cfg['ProjectionDataId']     = proj_id
    cfg['ReconstructionDataId'] = rec_id
    alg_id = astra.algorithm.create(cfg)
    astra.algorithm.run(alg_id)
    astra.algorithm.delete(alg_id)

In [ ]:
def test_reconstruction_cone_sirt(rec_volume,rec_volume_tmp, sinogram_volume, angles, proj_geom, gpus_list):
    astra.astra.set_gpu_index(gpus_list)
    angles = proj_geom['ProjectionAngles']
    

    
    angle_ids = np.arange(len(proj_geom['ProjectionAngles']))
    np.random.shuffle(angle_ids)
    ls =[1.5, 0.8,0.3]
    for l in ls:
        lsirt = l/sinogram_volume.shape[0]

        for angle_id in angle_ids:
            print angle_id,
            tt=time.time()
    #         print angle_id,
            tmp_proj_geom = proj_geom.copy()

            vol_geom = astra.create_vol_geom(rec_volume.shape)
            rec_id = astra.data3d.link('-vol', vol_geom, rec_volume)

            tmp_proj_geom['ProjectionAngles'] = np.array([tmp_proj_geom['ProjectionAngles'][angle_id],])

            plt.figure()
            plt.imshow(sinogram_volume[:, angle_id, :])
            plt.colorbar()
            plt.title('original projection')
            plt.show()

            proj_id = astra.data3d.create('-sino', tmp_proj_geom)
            performFP(proj_id, rec_id)

            tmp_projection = np.squeeze(astra.data3d.get(proj_id))
#             astra.data3d.delete(proj_id)
            astra.data3d.delete(rec_id)


        #     plt.figure()
        #     plt.imshow(tmp_projection)
        #     plt.colorbar()
        #     plt.title('tmp projection')
        #     plt.show()

            t = rec_volume[int(rec_volume.shape[0] / 2)].copy()

        #     plt.figure()
        #     plt.imshow(t)
        #     plt.title('source rec volume')
        #     plt.colorbar()
        #     plt.show()


            prj_diff =  tmp_projection - sinogram_volume[:, angle_id, :]

        #     plt.figure()
        #     plt.imshow(np.squeeze(prj_diff))
        #     plt.colorbar()
        #     plt.title('prj_diff')
        #     plt.show()

            prj_diff = -lsirt*np.require(prj_diff[:,np.newaxis,:],requirements=['C'])

            rec2_id = astra.data3d.link('-vol', vol_geom, rec_volume_tmp)
            proj2_id = astra.data3d.link('-sino', tmp_proj_geom, prj_diff)
            performBP(proj2_id, rec2_id)
#             astra.data3d.delete(proj2_id)
            astra.data3d.delete(rec2_id)
            rec_volume+=rec_volume_tmp
#             numexpr.evaluate('rec_volume+rec_volume_tmp',out=rec_volume)
        #     plt.figure()
        # #     plt.imshow(t- rec_volume[int(rec_volume.shape[0] / 2)])
        #     plt.imshow(rec_volume[int(rec_volume.shape[0] / 2)])
        #     plt.colorbar()
        #     plt.title('new vol')
        #     plt.show()
            print time.time()-tt

#     astra.data3d.info()
    astra.clear()

In [ ]:
plt.viridis()
t=time.time()
rec_volume,sinogram_volume, angles, proj_geom = test_reconstruction_cone_fdk(
    size=1200, angles_count= 50, gpus_list=[0,1])

print(time.time()-t)
plt.imshow(rec_volume[int(rec_volume.shape[0] / 2),:,:])
plt.colorbar()
plt.show()

# rec_volume[:]=0
rec_volume_tmp = np.memmap('/diskmnt/fast/makov/rec2.tmp',
                       dtype='float32', mode='w+',
                       shape=rec_volume.shape
                       )
t=time.time()
test_reconstruction_cone_sirt(rec_volume, rec_volume_tmp, sinogram_volume, angles, proj_geom, [0,1])
print(time.time()-t)
plt.imshow(rec_volume[int(rec_volume.shape[0] / 2),:,:])
plt.colorbar()
plt.show()

# plt.imshow(rec_volume_tmp[int(rec_volume_tmp.shape[0] / 2),:,:])
# plt.colorbar()
# plt.show()

In [ ]:
plt.imshow(rec_volume[int(rec_volume.shape[0] / 2),:,:])
plt.colorbar()
plt.show()

In [ ]:
plt.imshow(rec_volume_tmp[int(rec_volume.shape[0] / 2),:,:])
plt.colorbar()
plt.show()

In [ ]:
plt.imshow(sinogram_volume[50,:, :])

In [ ]:
sinogram_volume[:,:, 100][:,:,np.newaxis].shape

In [ ]:
sinogram_volume.shape

In [ ]:
astra.clear()

In [ ]:
pixel_size = (2.82473e-3)*8    
os_distance = 56.135 / pixel_size  # object-sample distance    
ds_distance = 225.082 / pixel_size  # detector-sample distance

print(os_distance)
print(ds_distance-os_distance)

In [ ]:
ds_distance / os_distance

In [ ]: