text: import os
  import numpy as np
  import time
  import sys
  
  import torch
  import torch.nn as nn
  import torch.backends.cudnn as cudnn
  import torchvision
  import torchvision.transforms as transforms
  import torch.optim as optim
  import torch.nn.functional as tfunc
  from torch.utils.data import DataLoader
  from torch.optim.lr_scheduler import ReduceLROnPlateau
  import torch.nn.functional as func
  
  from sklearn.metrics.ranking import roc_auc_score
  
  from bayersite.HeatmapGenerator import *
  from bayersite.DensenetModels import DenseNet121
  from bayersite.DensenetModels import DenseNet169
  from bayersite.DensenetModels import DenseNet201
  from bayersite.DatasetGenerator import DatasetGenerator
  
  
  #-------------------------------------------------------------------------------- 
  
  class Zorg ():
  
      #---- Train the densenet network 
      #---- pathDirData - path to the directory that contains images
      #---- pathFileTrain - path to the file that contains image paths and label pairs (training set)
      #---- pathFileVal - path to the file that contains image path and label pairs (validation set)
      #---- nnArchitecture - model architecture 'DENSE-NET-121', 'DENSE-NET-169' or 'DENSE-NET-201'
      #---- nnIsTrained - if True, uses pre-trained version of the network (pre-trained on imagenet)
      #---- nnClassCount - number of output classes 
      #---- trBatchSize - batch size
      #---- trMaxEpoch - number of epochs
      #---- transResize - size of the image to scale down to (not used in current implementation)
      #---- transCrop - size of the cropped image 
      #---- launchTimestamp - date/time, used to assign unique name for the checkpoint file
      #---- checkpoint - if not None loads the model and continues training
      
      def train (pathDirData, pathFileTrain, pathFileVal, nnArchitecture, nnIsTrained, nnClassCount, trBatchSize, trMaxEpoch, transResize, transCrop, launchTimestamp, checkpoint):
  
          
          #-------------------- SETTINGS: NETWORK ARCHITECTURE
          if nnArchitecture == 'DENSE-NET-121': model = DenseNet121(nnClassCount, nnIsTrained).cuda()
          elif nnArchitecture == 'DENSE-NET-169': model = DenseNet169(nnClassCount, nnIsTrained).cuda()
          elif nnArchitecture == 'DENSE-NET-201': model = DenseNet201(nnClassCount, nnIsTrained).cuda()
          
          model = torch.nn.DataParallel(model).cuda()
                  
          #-------------------- SETTINGS: DATA TRANSFORMS
          normalize = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
          
          transformList = []
          transformList.append(transforms.RandomResizedCrop(transCrop))
          transformList.append(transforms.RandomHorizontalFlip())
          transformList.append(transforms.ToTensor())
          transformList.append(normalize)      
          transformSequence=transforms.Compose(transformList)
  
          #-------------------- SETTINGS: DATASET BUILDERS
          datasetTrain = DatasetGenerator(pathImageDirectory=pathDirData, pathDatasetFile=pathFileTrain, transform=transformSequence)
          datasetVal =   DatasetGenerator(pathImageDirectory=pathDirData, pathDatasetFile=pathFileVal, transform=transformSequence)
                
          dataLoaderTrain = DataLoader(dataset=datasetTrain, batch_size=trBatchSize, shuffle=True,  num_workers=24, pin_memory=True)
          dataLoaderVal = DataLoader(dataset=datasetVal, batch_size=trBatchSize, shuffle=False, num_workers=24, pin_memory=True)
          
          #-------------------- SETTINGS: OPTIMIZER & SCHEDULER
          optimizer = optim.Adam (model.parameters(), lr=0.0001, betas=(0.9, 0.999), eps=1e-08, weight_decay=1e-5)
          scheduler = ReduceLROnPlateau(optimizer, factor = 0.1, patience = 5, mode = 'min')
                  
          #-------------------- SETTINGS: LOSS
          loss = torch.nn.BCELoss(size_average = True)
          
          #---- Load checkpoint 
          if checkpoint != None:
              modelCheckpoint = torch.load(checkpoint)
              model.load_state_dict(modelCheckpoint['state_dict'])
              optimizer.load_state_dict(modelCheckpoint['optimizer'])
  
          
          #---- TRAIN THE NETWORK
          
          lossMIN = 100000
          
          for epochID in range (0, trMaxEpoch):
              
              timestampTime = time.strftime("%H%M%S")
              timestampDate = time.strftime("%d%m%Y")
              timestampSTART = timestampDate + '-' + timestampTime
                           
              ChexnetTrainer.epochTrain (model, dataLoaderTrain, optimizer, scheduler, trMaxEpoch, nnClassCount, loss)
              lossVal, losstensor = ChexnetTrainer.epochVal (model, dataLoaderVal, optimizer, scheduler, trMaxEpoch, nnClassCount, loss)
              
              timestampTime = time.strftime("%H%M%S")
              timestampDate = time.strftime("%d%m%Y")
              timestampEND = timestampDate + '-' + timestampTime
              
              scheduler.step(losstensor.data[0])
              
              if lossVal < lossMIN:
                  lossMIN = lossVal    
                  torch.save({'epoch': epochID + 1, 'state_dict': model.state_dict(), 'best_loss': lossMIN, 'optimizer' : optimizer.state_dict()}, 'm-' + launchTimestamp + '.pth.tar')
                  print ('Epoch [' + str(epochID + 1) + '] [save] [' + timestampEND + '] loss= ' + str(lossVal))
              else:
                  print ('Epoch [' + str(epochID + 1) + '] [----] [' + timestampEND + '] loss= ' + str(lossVal))
                       
      #-------------------------------------------------------------------------------- 
         
      def epochTrain (model, dataLoader, optimizer, scheduler, epochMax, classCount, loss):
          
          model.train()
          
          for batchID, (input, target) in enumerate (dataLoader):
                          
              target = target.cuda(async = True)
                   
              varInput = torch.autograd.Variable(input)
              varTarget = torch.autograd.Variable(target)         
              varOutput = model(varInput)
              
              lossvalue = loss(varOutput, varTarget)
                         
              optimizer.zero_grad()
              lossvalue.backward()
              optimizer.step()
              
      #-------------------------------------------------------------------------------- 
          
      def epochVal (model, dataLoader, optimizer, scheduler, epochMax, classCount, loss):
          
          model.eval ()
          
          lossVal = 0
          lossValNorm = 0
          
          losstensorMean = 0
          
          for i, (input, target) in enumerate (dataLoader):
              
              target = target.cuda(async=True)
                   
              varInput = torch.autograd.Variable(input, volatile=True)
              varTarget = torch.autograd.Variable(target, volatile=True)    
              varOutput = model(varInput)
              
              losstensor = loss(varOutput, varTarget)
              losstensorMean += losstensor
              
              lossVal += losstensor.data[0]
              lossValNorm += 1
              
          outLoss = lossVal / lossValNorm
          losstensorMean = losstensorMean / lossValNorm
          
          return outLoss, losstensorMean
                 
      #--------------------------------------------------------------------------------     
       
      #---- Computes area under ROC curve 
      #---- dataGT - ground truth data
      #---- dataPRED - predicted data
      #---- classCount - number of classes
      
      def computeAUROC (dataGT, dataPRED, classCount):
          
          outAUROC = []
          
          datanpGT = dataGT.cpu().numpy()
          datanpPRED = dataPRED.cpu().numpy()
          
         # for i in range(classCount):
          #    outAUROC.append(roc_auc_score(datanpGT[:, i], datanpPRED[:, i]))
              
          return outAUROC
          
          
      #--------------------------------------------------------------------------------  
      
      #---- Test the trained network 
      #---- pathDirData - path to the directory that contains images
      #---- pathFileTrain - path to the file that contains image paths and label pairs (training set)
      #---- pathFileVal - path to the file that contains image path and label pairs (validation set)
      #---- nnArchitecture - model architecture 'DENSE-NET-121', 'DENSE-NET-169' or 'DENSE-NET-201'
      #---- nnIsTrained - if True, uses pre-trained version of the network (pre-trained on imagenet)
      #---- nnClassCount - number of output classes 
      #---- trBatchSize - batch size
      #---- trMaxEpoch - number of epochs
      #---- transResize - size of the image to scale down to (not used in current implementation)
      #---- transCrop - size of the cropped image 
      #---- launchTimestamp - date/time, used to assign unique name for the checkpoint file
      #---- checkpoint - if not None loads the model and continues training
      
      def test (self, pathDirData, pathFileTest, pathModel, nnArchitecture, nnClassCount, nnIsTrained, trBatchSize, transResize, transCrop, launchTimeStamp):   
          
          
          disease = [ 'Atelectasis', 'Cardiomegaly', 'Effusion', 'Infiltration', 'Mass', 'Nodule', 'Pneumonia',
                  'Pneumothorax', 'Consolidation', 'Edema', 'Emphysema', 'Fibrosis', 'Pleural_Thickening', 'Hernia']
          
          cudnn.benchmark = True
          
          #-------------------- SETTINGS: NETWORK ARCHITECTURE, MODEL LOAD
          if nnArchitecture == 'DENSE-NET-121': model = DenseNet121(nnClassCount, nnIsTrained).cuda()
          elif nnArchitecture == 'DENSE-NET-169': model = DenseNet169(nnClassCount, nnIsTrained).cuda()
          elif nnArchitecture == 'DENSE-NET-201': model = DenseNet201(nnClassCount, nnIsTrained).cuda()
          
          model = torch.nn.DataParallel(model).cuda() 
          
          modelCheckpoint = torch.load(pathModel)
          model.load_state_dict(modelCheckpoint['state_dict'])
  
          #-------------------- SETTINGS: DATA TRANSFORMS, TEN CROPS
          normalize = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
          
          #-------------------- SETTINGS: DATASET BUILDERS
          transformList = []
          transformList.append(transforms.Resize(transResize))
          transformList.append(transforms.TenCrop(transCrop))
          transformList.append(transforms.Lambda(lambda crops: torch.stack([transforms.ToTensor()(crop) for crop in crops])))
          transformList.append(transforms.Lambda(lambda crops: torch.stack([normalize(crop) for crop in crops])))
          transformSequence=transforms.Compose(transformList)
          
          datasetTest = DatasetGenerator(pathImageDirectory=pathDirData, pathDatasetFile=pathFileTest, transform=transformSequence)
          dataLoaderTest = DataLoader(dataset=datasetTest, batch_size=trBatchSize, num_workers=8, shuffle=False, pin_memory=True)
          generate(pathDirData +"/"+ datasetTest.get_path()) 
          outGT = torch.FloatTensor().cuda()
          outPRED = torch.FloatTensor().cuda()
          outMean = torch.FloatTensor().cuda()
          model.eval()
          dire = pathDirData
          num_files = len([x for x in os.listdir(dire) if os.path.isfile(os.path.join(dire,x))])
          print(num_files)
          for i, (input, target) in enumerate(dataLoaderTest):
              
              target = target.cuda()
              outGT = torch.cat((outGT, target), 0)
              
              bs, n_crops, c, h, w = input.size()
              
              varInput = torch.autograd.Variable(input.view(-1, c, h, w).cuda(), volatile=True)
              
              out = model(varInput)
              outMean = out.view(bs, n_crops, -1).mean(1)
              print(outMean*100)
              outPRED = torch.cat((outPRED, outMean.data), 0)
  
  
  
          #aurocIndividual = Zorg.computeAUROC(outGT, outPRED, nnClassCount)
          #aurocMean = np.array(aurocIndividual).mean()
          outMean = outMean*100
          #print ('AUROC mean ', aurocMean)
          print(len(outMean[0]))
          most_prob = 0;
          second_prob = 0;
          for i in range(len(outMean[0])):
              if(float(outMean[0][i]) > float(outMean[0][most_prob])):
                  second_prob = most_prob
                  most_prob = i
  
          output = open("diagnosis.txt", "w")
          drugs = ['Adalat', 'Adempas', 'Adengo', 'Afrin', 'Avelox,', 'Nexavar', 'Stivarga', 'Xofigo', 'Medrad', 'Medrad', 'Stivarga', 'Nexavar', 'Adempas', 'Adalat']
          if(float(outMean[0][most_prob] < 1*(10**-1))):
              print("Inconclusive")
              output.write("Inconclusive diagnosis, probabilities below threshold. Upload new photo or consult a physician. Perhaps you are fine!")
          else:
              print("Conclusion!")
              output.write("Primary Diagnosis:\nYou are most likely afflicted with: "+ disease[most_prob]+ "\nThis is concluded with a probability of: "+ str(float(outMean[0][most_prob])))
              output.write("\nSecondary Diagnosis:\nThe second most likely disease is: "+ disease[second_prob]+ "\nThis is concluded with a probability of: "+ str(float(outMean[0][second_prob])))
              output.write("\nAsk your doctor for: " + drugs[most_prob])
  
          for i in range (0, 14):
              print(disease[i], ' ', float(outMean[0][i]))
           
       
          return
          #print(int(outMean[0][1]))
  #-------------------------------------------------------------------------------- 
  
  
  
  
  
  