--- a
+++ b/drunet/segment.py
@@ -0,0 +1,320 @@
+import os
+import time
+import argparse
+import pathlib
+
+import tqdm
+import cv2 as cv
+import numpy as np
+import pandas as pd
+import tensorflow as tf
+from tensorflow import keras
+import matplotlib.pyplot as plt
+from tensorflow.compat.v1 import ConfigProto
+from tensorflow.compat.v1 import InteractiveSession
+# Custom package
+import data
+import loss
+import utils
+import module
+import performance
+from model import dr_unet
+
+config = ConfigProto()
+config.gpu_options.allow_growth = True
+
+# 1. Parameter settings
+parser = argparse.ArgumentParser(description="Segment Use Args")
+parser.add_argument('--model-name', default='DR_UNet', type=str)
+parser.add_argument('--dims', default=32, type=int)
+parser.add_argument('--epochs', default=50, type=int)
+parser.add_argument('--batch-size', default=16, type=int)
+parser.add_argument('--lr', default=2e-4, type=float)
+
+# Training data, testing, verification parameter settings
+parser.add_argument('--height', default=256, type=int)
+parser.add_argument('--width', default=256, type=int)
+parser.add_argument('--channel', default=1, type=int)
+parser.add_argument('--pred-height', default=4 * 256, type=int)
+parser.add_argument('--pred-width', default=4 * 256, type=int)
+parser.add_argument('--total-samples', default=5000, type=int)
+parser.add_argument('--invalid-samples', default=1000, type=int)
+parser.add_argument('--regularize', default=False, type=bool)
+parser.add_argument('--record-dir', default=r'', type=str, help='the save dir of tfrecord')
+parser.add_argument('--train-record-name', type=str, default=r'train_data', help='the train record save name')
+parser.add_argument('--test-image-dir', default=r'', type=str, help='the path of test images dir')
+parser.add_argument('--invalid-record-name', type=str, default=r'test_data', help='the invalid record save name')
+parser.add_argument('--gt-mask-dir', default=r'', type=str, help='the ground truth dir of validation set')
+parser.add_argument('--invalid-volume-dir', default=r'', type=str, help='estimation bleeding volume')
+args = parser.parse_args()
+
+
+class Segmentation:
+    def __init__(self, params):
+        self.params = params
+        self.input_shape = [params.height, params.width, params.channel]
+        self.mask_shape = [params.height, params.width, 1]
+        self.model_name = params.model_name
+        self.crop_height = params.pred_height
+        self.crop_width = params.pred_width
+        self.regularize = params.regularize
+
+        # Obtain a segmentation model
+        self.seg_model = dr_unet.dr_unet(input_shape=self.input_shape, dims=params.dims)
+        self.seg_model.summary()
+
+        # Optimization function
+        self.optimizer = tf.keras.optimizers.Adam(lr=params.lr)
+
+        # Every epoch, predict invalid-images to test the segmentation performance of the model
+        self.save_dir = str(params.model_name).upper()
+        self.weight_save_dir = os.path.join(self.save_dir, 'checkpoint')
+        self.pred_invalid_save_dir = os.path.join(self.save_dir, 'invalid_pred')
+        self.invalid_crop_save_dir = os.path.join(self.save_dir, 'invalid_pred_crop')
+        self.pred_test_save_dir = os.path.join(self.save_dir, 'test_pred')
+        utils.check_file([
+            self.save_dir, self.weight_save_dir, self.pred_invalid_save_dir,
+            self.pred_test_save_dir, self.invalid_crop_save_dir]
+        )
+
+        # Save model parameters
+        train_steps = tf.Variable(0, tf.int32)
+        self.save_ckpt = tf.train.Checkpoint(
+            train_steps=train_steps, seg_model=self.seg_model, model_optimizer=self.optimizer)
+        self.save_manger = tf.train.CheckpointManager(
+            self.save_ckpt, directory=self.weight_save_dir, max_to_keep=1)
+
+        # Set the loss function
+        self.loss_fun = loss.bce_dice_loss
+
+    def load_model(self):
+        if self.save_manger.latest_checkpoint:
+            self.save_ckpt.restore(self.save_manger.latest_checkpoint)
+            print('Loading model: {}'.format(self.save_manger.latest_checkpoint))
+        else:
+            print('Retrain the model！')
+        return
+
+    @tf.function
+    def train_step(self, inputs, target):
+        tf.keras.backend.set_learning_phase(True)
+
+        with tf.GradientTape() as tape:
+            pred_mask = self.seg_model(inputs)
+            loss = self.loss_fun(target, pred_mask)
+            if self.regularize:
+                loss = tf.reduce_sum(loss) + tf.reduce_sum(self.seg_model.losses)
+        gradient = tape.gradient(loss, self.seg_model.trainable_variables)
+        self.optimizer.apply_gradients(zip(gradient, self.seg_model.trainable_variables))
+        return tf.reduce_mean(loss)
+
+    @tf.function
+    def inference(self, inputs):
+        tf.keras.backend.set_learning_phase(True)
+        pred = self.seg_model(inputs)
+        return pred
+
+    @staticmethod
+    def calculate_volume_by_mask(mask_dir, save_dir, model_name, dpi=96, thickness=0.45):
+        all_mask_file_paths = utils.list_file(mask_dir)
+
+        pd_record = pd.DataFrame(columns=['file_name', 'Volume'])
+        for file_dir in tqdm.tqdm(all_mask_file_paths):
+            file_name = pathlib.Path(file_dir).stem
+
+            each_blood_volume = module.calculate_volume(file_dir, thickness=thickness, dpi=dpi)
+            pd_record = pd_record.append({'file_name': file_name, 'Volume': each_blood_volume}, ignore_index=True)
+            pd_record.to_csv(
+                os.path.join(save_dir, '{}_{}.csv'.format(model_name, file_name)), index=True, header=True)
+        return
+
+    def predict_blood_volume(self, input_dir, save_dir, calc_nums=-1, dpi=96, thickness=0.45):
+        """
+         :param input_dir: The directory for testing bleeding volume images,
+                           there are multiple folders under the directory, each folder represents a CT image of a patient
+         :param save_dir: The predicted segmented image save directory
+         :param calc_nums: predict how many images in the folder
+         :param dpi: image parameters
+         :param thickness: slice thickness
+        """
+        # Loading weights of model
+        self.load_model()
+        save_pred_images_dir = os.path.join(save_dir, 'pred_images')
+        save_pred_csv_dir = os.path.join(save_dir, 'pred_csv')
+        utils.check_file([save_pred_images_dir, save_pred_csv_dir])
+        all_file_dirs = utils.list_file(input_dir)
+
+        cost_time_list = []
+        total_images = 0
+        for file_dir in tqdm.tqdm(all_file_dirs):
+            file_name = pathlib.Path(file_dir).stem
+
+            image_names, ori_images, normed_images = data.get_test_data(
+                test_data_path=file_dir, image_shape=self.input_shape, image_nums=calc_nums)
+            total_images += len(image_names)
+
+            start_time = time.time()
+            pred_mask = self.inference(normed_images)
+            end_time = time.time()
+            print('FPS: {}'.format(pred_mask.shape[0] / (end_time - start_time)),
+                  pred_mask.shape, end_time - start_time)
+
+            denorm_pred_mask = module.reverse_pred_image(pred_mask.numpy())  # (image_nums, 256, 256, 1)
+            if denorm_pred_mask.ndim == 2 and self.input_shape[-1] == 1:
+                denorm_pred_mask = np.expand_dims(denorm_pred_mask, 0)
+
+            drawed_images = []
+            blood_areas = []
+            pd_record = pd.DataFrame(columns=['image_name', 'Square Centimeter', 'Volume'])
+            for index in range(denorm_pred_mask.shape[0]):
+                drawed_image, blood_area = module.draw_contours(ori_images[index], denorm_pred_mask[index], dpi=dpi)
+                drawed_images.append(drawed_image)
+                blood_areas.append(blood_area)
+                pd_record = pd_record.append({'image_name': image_names[index], 'Square Centimeter': blood_area},
+                                             ignore_index=True)
+
+            one_pred_save_dir = os.path.join(save_pred_images_dir, file_name)
+            module.save_invalid_data(ori_images, drawed_images, denorm_pred_mask,
+                                     image_names, reshape=True, save_dir=one_pred_save_dir)
+
+            # Calculate the amount of bleeding based on the area of each layer of hematoma
+            blood_volume = module.count_volume(blood_areas, thickness=thickness)
+            pd_record = pd_record.append({'Volume': blood_volume}, ignore_index=True)
+            pd_record.to_csv(os.path.join(save_pred_csv_dir, '{}_{}.csv'.format(self.model_name, file_name)),
+                             index=True, header=True)
+            cost_time_list.append(end_time - start_time)
+            print('FileName: {} time: {}'.format(file_name, end_time - start_time))
+        print('total_time: {:.2f}, mean_time: {:.2f}, total_images: {}'.format(
+            np.sum(cost_time_list), np.mean(cost_time_list), total_images))
+        return
+
+    def predict_and_save(self, input_dir, save_dir, calc_nums=-1, batch_size=16):
+        """ predict bleeding image and save
+         :param input_dir: There are several images waiting to be tested under the input_dir folder
+         :param save_dir: The file directory where the segmented image predicted by the model is saved
+         :param calc_nums: How many images are taken from the directory to participate in the calculation
+         :param batch_size: how many images to test each time
+         :return:
+        """
+        mask_save_dir = os.path.join(save_dir, 'pred_mask')
+        drawed_save_dir = os.path.join(save_dir, 'drawed_image')
+        utils.check_file([mask_save_dir, drawed_save_dir])
+        self.load_model()
+
+        test_image_list = utils.list_file(input_dir)
+        for index in range(len(test_image_list) // 128 + 1):
+            input_test_list = test_image_list[index * 128:(index + 1) * 128]
+
+            image_names, ori_images, normed_images = data.get_test_data(
+                test_data_path=input_test_list, image_shape=self.input_shape, image_nums=-1,
+            )
+            if calc_nums != -1:
+                ori_images = ori_images[:calc_nums]
+                normed_images = normed_images[:calc_nums]
+                image_names = image_names[:calc_nums]
+
+            inference_times = normed_images.shape[0] // batch_size + 1
+            for inference_time in range(inference_times):
+                this_normed_images = normed_images[
+                                     inference_time * batch_size:(inference_time + 1) * batch_size, ...]
+                this_ori_images = ori_images[
+                                  inference_time * batch_size:(inference_time + 1) * batch_size, ...]
+                this_image_names = image_names[
+                                   inference_time * batch_size:(inference_time + 1) * batch_size]
+
+                this_pred_mask = self.inference(this_normed_images)
+                this_denorm_pred_mask = module.reverse_pred_image(this_pred_mask.numpy())
+                if ori_images.shape[0] == 1:
+                    this_denorm_pred_mask = np.expand_dims(this_denorm_pred_mask, 0)
+
+                for i in range(this_denorm_pred_mask.shape[0]):
+                    bin_denorm_pred_mask = this_denorm_pred_mask[i]
+                    this_drawed_image, this_blood_area = module.draw_contours(
+                        this_ori_images[i], bin_denorm_pred_mask, dpi=96
+                    )
+                    cv.imwrite(os.path.join(
+                        mask_save_dir, '{}'.format(this_image_names[i])), bin_denorm_pred_mask
+                    )
+                    cv.imwrite(os.path.join(
+                        drawed_save_dir, '{}'.format(this_image_names[i])), this_drawed_image
+                    )
+        return
+
+    def train(self, start_epoch=1):
+        # get training dataset
+        train_data = data.get_tfrecord_data(
+            self.params.record_dir, self.params.train_record_name,
+            self.input_shape, batch_size=self.params.batch_size)
+        self.load_model()
+
+        pd_record = pd.DataFrame(columns=['Epoch', 'Iteration', 'Loss', 'Time'])
+        data_name, original_test_image, norm_test_image = data.get_test_data(
+            test_data_path=self.params.test_image_dir, image_shape=self.input_shape, image_nums=-1
+        )
+
+        start_time = time.time()
+        best_dice = 0.0
+        for epoch in range(start_epoch, self.params.epochs):
+            for train_image, gt_mask in tqdm.tqdm(
+                    train_data, total=self.params.total_samples // self.params.batch_size):
+                self.save_ckpt.train_steps.assign_add(1)
+                iteration = self.save_ckpt.train_steps.numpy()
+
+                # training step
+                train_loss = self.train_step(train_image, gt_mask)
+                if iteration % 100 == 0:
+                    print('Epoch: {}, Iteration: {}, Loss: {:.2f}, Time: {:.2f} s'.format(
+                        epoch, iteration, train_loss, time.time() - start_time))
+
+                    # test step
+                    test_pred = self.inference(norm_test_image)
+                    module.save_images(
+                        image_shape=self.mask_shape, pred=test_pred,
+                        save_path=self.pred_test_save_dir, index=iteration, split=False
+                    )
+                    pd_record = pd_record.append({
+                        'Epoch': epoch, 'Iteration': iteration, 'Loss': train_loss.numpy(),
+                        'Time': time.time() - start_time}, ignore_index=True
+                    )
+                    pd_record.to_csv(os.path.join(
+                        self.save_dir, '{}_record.csv'.format(self.params.model_name)), index=True, header=True
+                    )
+
+            m_dice = self.invalid(epoch)
+            if m_dice > best_dice:
+                best_dice = m_dice
+                print('Best Dice:{}'.format(best_dice))
+                self.save_manger.save(checkpoint_number=epoch)
+        return
+
+    def invalid(self, epoch):
+        invalid_data = data.get_tfrecord_data(
+            self.params.record_dir, self.params.invalid_record_name,
+            self.input_shape, batch_size=self.params.batch_size, shuffle=False)
+
+        epoch_pred_save_dir = None
+        for index, (invalid_image, invalid_mask) in enumerate(
+                tqdm.tqdm(invalid_data, total=self.params.invalid_samples // self.params.batch_size + 1)):
+            invalid_pred = self.inference(invalid_image)
+            epoch_pred_save_dir = os.path.join(self.pred_invalid_save_dir, f'epoch_{epoch}')
+            module.save_images(
+                image_shape=self.mask_shape, pred=invalid_pred,
+                save_path=epoch_pred_save_dir, index=f'{index}', split=False
+            )
+
+        # Test model performance
+        epoch_cropped_save_dir = os.path.join(
+            self.invalid_crop_save_dir, f'epoch_{epoch}'
+        )
+        utils.crop_image(epoch_pred_save_dir, epoch_cropped_save_dir,
+                         self.crop_width, self.crop_height,
+                         self.input_shape[0], self.input_shape[1]
+                         )
+        m_dice, m_iou, m_precision, m_recall = performance.save_performace_to_csv(
+            pred_dir=epoch_cropped_save_dir, gt_dir=self.params.gt_mask_dir,
+            img_resize=(self.params.height, self.params.width),
+            csv_save_name=f'{self.model_name}_epoch_{epoch}',
+            csv_save_path=epoch_cropped_save_dir
+        )
+        return m_dice
+