<filename>Organ_Finder_NN_Train.py
import numpy as np
import sys

from PIL import ImageEnhance
from PIL import Image
from torchvision import transforms, datasets,utils

import cv2
import albumentations
from albumentations import Compose, RandomCrop, Normalize, HorizontalFlip, Resize,Transpose,ElasticTransform
from albumentations.pytorch.transforms import ToTensor
import albumentations as albu

from glob import *
import pandas as pd
import torch
import os
import re
import time
from tqdm import tqdm
import pydicom

torch.manual_seed(0)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = True

import torch.nn.functional as F

import torchvision
from os import listdir
from os.path import isfile, join

from itertools import chain
import itertools as it

from datetime import datetime, date

from ast import literal_eval

import time
from datetime import datetime

from random import shuffle

from Training_model import train_model
from getting_models import initialize_model

import warnings
import random

import torch.nn as nn
from utils.yaml_config_hook import yaml_config_hook
from torch.utils.data.sampler import WeightedRandomSampler
from sklearn.model_selection import KFold


def Location_Specifier(Appy_df,Type,List_location):
    """
    Arguments
    ----
    List_location: Should either be a string in a list, or multiple
    strings in a list specifying the region of interest
        It should be noted that the capitalization of the word
        does not matter
    Appy_df: All list of all the files in the path of interest

    Return
    ---
    Specific_df: Returns a cleaned dataframe of Appy_df, including
    only the paths with the specified location as found in the Google OCR
    results

    There is still a question of inference of images that do not have
    text found on the image, and whether we can infer a location label
    onto these images using the training of the Google OCR results
    """

    File_Dir= r'C:\Users\<NAME>\Desktop\Appendicitis'

    Original_location_word= List_location[0]
    Organ_Mapping= pd.read_excel(os.path.join(File_Dir,'Organ_mapping.xlsx'))

    Organ_Mapping = Organ_Mapping[Organ_Mapping['Conserved Term'].isin(List_location)]

    List_location= List_location+ Organ_Mapping['Word'].to_list()
    quoted = re.compile('"[^"]*"')
    for index in range(len(List_location)):
        List_location[index] = List_location[index].replace('"', '')
        List_location[index] = List_location[index].lower()
        List_location[index] = re.sub(r'[^\w\s]','',List_location[index])

    List_location= list(dict.fromkeys(List_location))
    print(List_location)

    Conserved_term = Organ_Mapping['Conserved Term'].to_list

    Appy_df= Appy_df[Appy_df['Type']== str(Type)]
    Appy_dict= {}

    if len(List_location) ==1:
        List_location.append('ONLY_ONE')
    for i in tqdm(range(len(Appy_df))):

        counter={}
        neg_counter=0

        for location in (List_location):
            if location != 'ONLY_ONE':

                row= Appy_df.iloc[i]
                try:
                    inter = [x.split() for x in row['OCR'].split(',')]
                    inter= re.split(",",row['OCR'])
                    inter= inter[1:]
                    Larger_list=[]
                    keep_word=[]
                    for word in inter:
                        quoted = re.compile('"[^"]*"')
                        for extr_word in quoted.findall(word):
                            clean_word = extr_word.replace('"', '')
                            clean_word= clean_word.lower()
                            clean_word= re.sub(r'[^\w\s]','',clean_word)
                            if clean_word =='cm':
                                neg_counter=1
                            if clean_word == str(location):

                                counter[clean_word]=1
                            for letter_index in range(len(clean_word)):
                                if letter_index+2 in range(len(clean_word)):
                                    if clean_word[letter_index] =='c'  and clean_word[letter_index+1] == 'm':
                                        neg_counter=1
                    if neg_counter ==0:
                        if List_location[0] in counter:
                            Appy_dict[row['Total_path']]= [row['Label'],row['OCR']]
                except ValueError:
                    print('Could not open path')

    Inbetween_df= pd.DataFrame(list(Appy_dict.items()),columns=[ 'Dicom_path','Nested'])
    Total_df= pd.concat([Inbetween_df.Dicom_path, Inbetween_df.Nested.apply(pd.Series) ], axis = 1)
    Total_df.columns = ['Dicom_path', 'Label', 'Location']
    return Total_df

def Model_wrapper(train,test,albumentations_transform,albumentations_transform_test,Pipeline_folder):
    warnings.filterwarnings("ignore", category=FutureWarning)

    config_file = os.path.join(Pipeline_folder,r'src\config.yml')
    config_dict = yaml_config_hook(config_file)

    criterion = ''

    batch_size = config_dict['batch_size']
    num_epochs = config_dict['epochs']
    lr = config_dict['learning_rate']
    Name_Models=['resnet50']

    k_folds = 5
    kfold = KFold(n_splits=k_folds, shuffle=True)

    dataset = train.append(pd.DataFrame(data = test), ignore_index=True)

    # K-fold Cross Validation model evaluation
    for fold, (train_ids, test_ids) in enumerate(kfold.split(dataset)):

        print(f'FOLD {fold}')
        print('--------------------------------')

        train = dataset.iloc[train_ids]
        test =  dataset.iloc[test_ids]

        Data_images_train=UltrasoundDataset(train, transform=albumentations_transform)
        Data_images_test=UltrasoundDataset(test, transform=albumentations_transform_test)

        target=torch.tensor(np.array(train['Labels'].to_list())).type(torch.LongTensor)

        # Compute samples weight (each sample should get its own weight)
        class_sample_count = torch.tensor(
            [(target == t).sum() for t in torch.unique(target, sorted=True)])
        weight = 1. / class_sample_count.float()
        samples_weight = torch.tensor([weight[t] for t in target])

        sampler = WeightedRandomSampler(samples_weight, len(samples_weight))

        train_loader = torch.utils.data.DataLoader(Data_images_train, batch_size=batch_size, sampler=sampler)
        validation_loader = torch.utils.data.DataLoader(Data_images_test, batch_size=512,shuffle=True)

        model= initialize_model(Name_Models[0],5)

        # Putting it on the two GPUs
        model = torch.nn.DataParallel(model).cuda()

        if config_dict['optimizer'] == 'Adam':
            optimizer = torch.optim.Adam(model.parameters(),
                                        lr=lr,
                                        weight_decay=config_dict['weight_decay'])

        if not os.path.exists(os.path.join(Pipeline_folder,'4_13_21',str(fold),'Trained_models',Name_Models[0])):
            os.makedirs(os.path.join(Pipeline_folder,'4_13_21',str(fold),'Trained_models',Name_Models[0]))
        test.to_csv(os.path.join(Pipeline_folder,'4_13_21',str(fold),'Trained_models','test.csv'))

        print(config_dict)

        model_trained= train_model(model,Name_Models[0],train_loader,validation_loader, criterion,
                                                optimizer,num_epochs,Pipeline_folder,5,config_dict)

class UltrasoundDataset():
    def __init__(self,Appy_df,transform):
        self.Appy_df= Appy_df
        self.transform = transform

    def __getitem__(self, index):
        path= self.Appy_df['TOTAL_PATH'].iloc[index]
        dataset = pydicom.dcmread(os.path.normpath(path),force=True)
        image_np=  dataset.pixel_array

        img = Image.fromarray(image_np)
        if hasattr(dataset, 'SequenceOfUltrasoundRegions'):
            region= dataset.data_element('SequenceOfUltrasoundRegions').value[0]
            try:
                region= dataset.data_element('SequenceOfUltrasoundRegions').value[1]
            except IndexError:
                pass

            img = img.crop((region.RegionLocationMinX0, region.RegionLocationMinY0,
                              region.RegionLocationMaxX1, region.RegionLocationMaxY1))
            image_np= np.array(img)

            height= image_np.shape[0]
            width= image_np.shape[1]

            pad_trans= albu.CenterCrop(int(np.ceil(height/1.2)),int(np.ceil(width/2)))

            augmented = pad_trans(image=(image_np))
            image_np = augmented['image']

            img = Image.fromarray(image_np)

        else:
            width, height = img.size
            new_width= 600
            new_height= 600
            left = (width - new_width)/2
            top = (height - new_height)/2
            right = (width + new_width)/2
            bottom = (height + new_height)/2
            bottom = (height + new_height)/2

            # Crop the center of the image
            img = img.crop((left, top, right, bottom))

        img= img.convert("L")
        image_np= np.array(img)
        image_np = np.repeat(image_np[..., np.newaxis], 3, -1)

        augmented = self.transform(image=(image_np))
        image = augmented['image']

        finding= self.Appy_df['Labels'].iloc[index]
        return (path,image, int(finding))

    def __len__(self):
        return len(self.Appy_df)


def main():

    Organ_Finder = True
    K_steps = True
    Label_change=True
    Only_healthy = False
    Only_disease_RLQ=False
    split_value=.7
    Pipeline_folder= r''

    # Training Transformations
    albumentations_transform = Compose([
    albu.Resize(224, 224,always_apply=True, p=1),
    albu.CoarseDropout(p=.3),
    albu.HorizontalFlip(p=0.4),
    albu.Rotate(p=.4),
    Normalize(
    mean=[0.485, 0.456, 0.406],
    std=[0.229, 0.224, 0.225],  p=1.0),
    ToTensor()
    ])

    # Testing Transformations
    albumentations_transform_test = Compose([
    Resize(224, 224,always_apply=True, p=1, interpolation=cv2.INTER_CUBIC),
    Normalize(
    mean=[0.485, 0.456, 0.406],
    std=[0.229, 0.224, 0.225],  p=1.0),
    ToTensor()])

    if os.path.exists(os.path.join(Pipeline_folder,'CLEAN_test.csv')):
        train = pd.read_csv(os.path.join(Pipeline_folder,'train.csv'))
        test = pd.read_csv(os.path.join(Pipeline_folder,'CLEAN_test.csv'))

    if Organ_Finder:
        print('Starting Training of Organ Finder')

        train['TOTAL_PATH'] = train['TOTAL_PATH'].apply(lambda x: os.path.normpath(x))
        test['TOTAL_PATH'] = test['TOTAL_PATH'].apply(lambda x: os.path.normpath(x))

        train = train[np.isfinite(train['Labels'])]
        print('Train: Number of Patients:',len(train["MRN"].unique()))
        print('Train: Number of Images:',len(train))


        test = test[np.isfinite(test['Labels'])]
        print('Test: Number of Patients:',len(test["MRN"].unique()))
        print('test: Number of Images:',len(test))


        if Only_healthy:

            # Getting rid of all RLQ in Appy patients
            print(train)
            APPY_train = train[train.LABEL_y == 'Appendicitis']
            APPY_train = APPY_train[APPY_train.LABEL_y != 4]

            # Untouched
            No_APPY_train = train[train.LABEL_y == 'No Appendicitis']
            train= pd.concat([APPY_train, No_APPY_train])

        if Only_disease_RLQ:

            # Untouched
            APPY_train = train[train.LABEL_y == 'Appendicitis']

            # Getting rid of all RLQ in Appy patients
            No_APPY_train = train[train.LABEL_y == 'No Appendicitis']
            No_APPY_train = No_APPY_train[No_APPY_train.Labels != 4]

            train= pd.concat([APPY_train, No_APPY_train])


    if K_steps:
      if os.path.exists(os.path.join(Pipeline_folder,'K_STEPS_CLEAN_test.csv')):
        train = pd.read_csv(os.path.join(Pipeline_folder,'K_STEPS_CLEAN_train.csv'))
        test = pd.read_csv(os.path.join(Pipeline_folder,'K_STEPS_CLEAN_test.csv'))
      else:
        good_files=[]
        for index in range(len(test)):
            try:
                dataset = pydicom.dcmread(test['TOTAL_PATH'].iloc[index],force=True)
                good_files.append(test['TOTAL_PATH'].iloc[index])

            except:
                pass # doing nothing on exception
        test= test[test ['TOTAL_PATH'].isin(good_files)]
        test.to_csv(os.path.join(Pipeline_folder,'K_STEPS_CLEAN_test.csv'))

        good_files=[]
        for index in range(len(train)):
            try:
                dataset = pydicom.dcmread(train['TOTAL_PATH'].iloc[index],force=True)
                good_files.append(train['TOTAL_PATH'].iloc[index])

            except:
                pass # doing nothing on exception
        train= train[train ['TOTAL_PATH'].isin(good_files)]
        train.to_csv(os.path.join(Pipeline_folder,'K_STEPS_CLEAN_train.csv'))

    print("----------------------------")
    print(len(train[train['Labels'] == 0]))
    print(len(train[train['Labels'] == 1]))
    print(len(train[train['Labels'] == 2]))
    print(len(train[train['Labels'] == 3]))
    print(len(train[train['Labels'] == 4]))
    print("----------------------------")
    print('Path Destination:',Pipeline_folder)

    Model_wrapper(train,test,albumentations_transform,albumentations_transform_test,Pipeline_folder)

if __name__ == '__main__':
    main()
