import tensorflow as tf
import os
from model import CycleGAN
import utils
import numpy as np
import nibabel as nib
import SimpleITK as sitk

FLAGS = tf.flags.FLAGS
tf.flags.DEFINE_string('model', 'pretrained/cbct2sct.pb', 'model path (.pb)')
tf.flags.DEFINE_integer('height', '288', 'height, default: 288')
tf.flags.DEFINE_integer('width', '320', 'width, default: 320')
# tf.flags.DEFINE_integer('depth_size', '64', 'image size, default: 64')


cbct_list = sorted(os.listdir('cbct/'))


def inference():
    for i in range(len(cbct_list)):
        filename = 'cbct/' + cbct_list[i]
        # cbct = sitk.ReadImage(filename)
        # cbct_data = sitk.GetArrayFromImage(cbct)
        cbct = nib.load(filename)
        cbct_data = cbct.get_fdata().swapaxes(0, 2)
        cbct_data[cbct_data > 3072] = 3072
        depth = cbct_data.shape[0]
        output_sample = np.zeros((depth, FLAGS.height, FLAGS.width))
        for i in range(depth):
            new_cbct = cbct_data[i, :, :].reshape(FLAGS.height, FLAGS.width, 1)
            new_cbct = new_cbct.astype(np.float32)
            input_image = utils.convert2float(new_cbct)
            graph = tf.Graph()

            with graph.as_default():
                with tf.gfile.FastGFile(FLAGS.model, 'rb') as model_file:
                    graph_def = tf.GraphDef()
                    graph_def.ParseFromString(model_file.read())
                [output_image] = tf.import_graph_def(graph_def, input_map={'input_image': input_image},
                                                     return_elements=['output_image:0'], name='output')
            print('12345')
            with tf.Session(graph=graph) as sess:
                generated = output_image.eval()
                generated = generated.reshape(FLAGS.height, FLAGS.width)
                output_sample[i, :, :] = np.round(generated, 0)
        # out = sitk.GetImageFromArray(output_sample)
        out = nib.Nifti1Image(output_sample.swapaxes(0, 2), cbct.affine)
        # sitk.WriteImage(out, 'sct_' + os.path.basename(filename))
        nib.save(out, 'sct_' + os.path.basename(filename))


def main(unused_argv):
    inference()


if __name__ == '__main__':
    tf.app.run()
