

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
    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,
                                       detector_size, detector_size, 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(size, angles_count, 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 = np.linspace(0, np.pi, angles_count, False)
#     angles = angles-(angles.max()+angles.min())/2

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

    # Create projection geometry
    proj_geom = build_geometry(angles, cube.shape[0])

    # Create volume desscription
    vol_geom = astra.create_vol_geom(cube.shape)

    # Upload data to ASTRA
    cube_id = astra.data3d.link('-vol', vol_geom, cube)

    print(time.time() - t)
    print('Creating sinogram')
    t = time.time()
    # 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(time.time() - t)
    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['GridRowCount'],   # TODO: check for non cube
                                  vol_geom['GridColCount'],
                                  vol_geom['GridSliceCount'])
                           )
    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('CGLS3D_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, 100)

    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