# code.py
# This is an image preprocessing code for HandBoneXRay data, please adopt env.yaml to
# create conda env to run this script
#
# author  : deng
# date    : 20230919
# platform: MacBook Pro 14 2021

import os
from shutil import rmtree

import cv2
import pydicom
import numpy as np
import matplotlib.pyplot as plt


if __name__ == '__main__':

    xray_dir = 'HandBoneXRay/'
    result_dir = 'results'
    plot = False

    # Create result dir to save processed images
    if os.path.isdir(result_dir):
        rmtree(result_dir)
    os.makedirs(result_dir)

    # Get xray dicom paths
    dicom_paths = [os.path.join(xray_dir, file_name)
                  for file_name in os.listdir(xray_dir)
                  if file_name.endswith('dcm')]

    for dicom_path in dicom_paths:
        print(f'Start to process {dicom_path}')

        # Read image
        ds = pydicom.dcmread(dicom_path)
        pixel_array = ds.pixel_array
        orig_pixel_array = pixel_array.copy()
        if plot:
            plt.title('Raw image')
            plt.imshow(pixel_array, cmap='gray')
            plt.show()

        # Thresholding
        pixel_array[pixel_array < 1] = 0
        pixel_array[pixel_array >= 1] = 1
        pixel_array = pixel_array.astype(np.uint8)
        if plot:
            plt.title('Thresholding')
            plt.imshow(pixel_array, cmap='gray')
            plt.show()

        # Calculate rotation angle by minAreaRect method
        contours, _ = cv2.findContours(pixel_array, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        contour_w_max_area = max(contours, key = cv2.contourArea)
        rect = cv2.minAreaRect(contour_w_max_area)
        angle = rect[2]
        if angle > 45:
            angle -= 90

        # Rotate image
        h, w = orig_pixel_array.shape[:2]
        center = rect[0]
        matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
        rotated = cv2.warpAffine(orig_pixel_array, matrix, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_CONSTANT)
        if plot:
            plt.title('Rotated image')
            plt.imshow(rotated, cmap='gray')
            plt.show()

        # Crop edge
        xs, ys = np.nonzero(rotated)
        x_min,x_max = xs.min(), xs.max()
        y_min, y_max = ys.min(), ys.max()
        cropped = rotated[x_min:x_max+1, y_min:y_max+1]
        cropped = cropped.astype(np.uint16)
        if plot:
            plt.title('Cropped image')
            plt.imshow(cropped, cmap='gray')
            plt.show()

        # Save it
        save_path = os.path.join(result_dir, os.path.basename(dicom_path))
        ds.PixelData = cropped
        ds.Rows = cropped.shape[0]
        ds.Columns = cropped.shape[1]
        pydicom.filewriter.dcmwrite(save_path, ds)
        if plot:
            ds = pydicom.dcmread(save_path)
            pixel_array = ds.pixel_array
            plt.title('Saved image')
            plt.imshow(pixel_array, cmap='gray')
            plt.show()

        print('done.')