#!/usr/bin/env python
# coding: utf-8

# In[1]:


#!/usr/bin/env python
# coding: utf-8

# # Detecting COVID-19 with Chest X Ray using PyTorch
# 
# Image classification of Chest X Rays in one of three classes: Normal, Viral Pneumonia, COVID-19

# # Importing Libraries

# In[28]:



import os
import shutil
import random
import torch
import torchvision
import numpy as np


from IPython import get_ipython
ipython_shell = get_ipython()
get_ipython().run_line_magic('matplotlib', 'inline')
from PIL import Image
from matplotlib import pyplot as plt

torch.manual_seed(0)

print('Using PyTorch version', torch.__version__)


# # Preparing Training and Test Sets

# In[2]:


class_names = ['normal', 'viral', 'covid']
root_dir = 'COVID-19 Radiography Database'
source_dirs = ['NORMAL', 'Viral Pneumonia', 'COVID-19']

if os.path.isdir(os.path.join(root_dir, source_dirs[1])):
	os.mkdir(os.path.join(root_dir, 'test'))

	for i, d in enumerate(source_dirs):
		os.rename(os.path.join(root_dir, d), os.path.join(root_dir, class_names[i]))

	for c in class_names:
		os.mkdir(os.path.join(root_dir, 'test', c))

	for c in class_names:
		images = [x for x in os.listdir(os.path.join(root_dir, c)) if x.lower().endswith('png')]
		selected_images = random.sample(images, 30)
		for image in selected_images:
			source_path = os.path.join(root_dir, c, image)
			target_path = os.path.join(root_dir, 'test', c, image)
			shutil.move(source_path, target_path)


# # Creating Custom Dataset

# In[60]:


class ChestXrayDataset(torch.utils.data.Dataset):
	def __init__(self,image_dirs,transform):
		def get_images(class_name):
			images=[x for x in os.listdir(image_dirs[class_name])if x.lower().endswith('png')]
			print(f'Found{len(images)}{class_name}examples')
			return images
		
		self.images={}
		self.class_names=['normal','viral','covid']
		
		for c in self .class_names:
			self.images[c]=get_images(c)
			
		self.image_dirs=image_dirs
		self.transform=transform
		
	def __len__(self):
		return sum([len(self.images[c])for c in self.class_names])
	
	#Most Important function

	def __getitem__(self,index):
		class_name=random.choice(self.class_names)
		index=index%len(self.images[class_name])  #To prevent Out of bounds
		image_name=self.images[class_name][index]
		image_path=os.path.join(self.image_dirs[class_name],image_name)
		image=Image.open(image_path).convert('RGB') #Not needed for CNN if its coded from start
		return self.transform(image),self.class_names.index(class_name)
		#Tensor Transformation in above line


# # Image Transformations

# In[61]:


train_transform=torchvision.transforms.Compose([
	torchvision.transforms.Resize(size=(224,224)), #Augmentation purpose 1
	torchvision.transforms.RandomHorizontalFlip(), #Augmentation purpose 2
	torchvision.transforms.ToTensor(),#Tensor conversion
	torchvision.transforms.Normalize(mean=[0.485,0.456,0.406],
									std=[0.229,0.224,0.225]) #Because Normalization done on Imagenet on Resnet18
])
#Same for test transformation except augmentation


# In[62]:


test_transform=torchvision.transforms.Compose([
	torchvision.transforms.Resize(size=(224,224)), #Augmentation purpose 1
	#REMOVED torchvision.transforms.RandomHorizontalFlip(), #Augmentation purpose 2
	torchvision.transforms.ToTensor(),#Tensor conversion
	torchvision.transforms.Normalize(mean=[0.485,0.456,0.406],
									std=[0.229,0.224,0.225]) #Because Normalization done on Imagenet on Resnet18
])
#Same for test transformation except augmentation 2


# # Prepare DataLoader

# 

# In[63]:


train_dirs = {
	'normal': 'COVID-19 Radiography Database/normal',
	'viral': 'COVID-19 Radiography Database/viral',
	'covid': 'COVID-19 Radiography Database/covid'
}

train_dataset=ChestXrayDataset(train_dirs,train_transform)


# In[64]:


test_dirs = {
	'normal': 'COVID-19 Radiography Database/test/normal',
	'viral': 'COVID-19 Radiography Database/test/viral',
	'covid': 'COVID-19 Radiography Database/test/covid'
}
test_dataset=ChestXrayDataset(test_dirs,test_transform)


# In[65]:


batch_size=6
dl_train=torch.utils.data.DataLoader(train_dataset,batch_size=batch_size,shuffle=True)
dl_test=torch.utils.data.DataLoader(test_dataset,batch_size=batch_size,shuffle=True)

print('No of Training Batches: ',len(dl_train))
print('No of Test batches:', len(dl_test))


# # Data Visualization

# In[66]:


class_names=train_dataset.class_names

def show_images(images,label,preds):
	plt.figure(figsize=(8,4))
	for i,image in enumerate(images):
		plt.subplot(1,6,i+1,xticks=[],yticks=[])
		image=image.numpy().transpose((1,2,0))#Converting from Tensor to numpy array
		mean=np.array([0.485,0.456,0.406])
		std=np.array([0.229,0.224,0.225])
		image=image*std+mean
		image=np.clip(image,0.,1.)
		plt.imshow(image)
		
		col='green' if preds[i]==labels[i] else 'red'
		
		plt.xlabel(f'{class_names[int(labels[i].numpy())]}')
		plt.ylabel(f'{class_names[int(preds[i].numpy())]}',color=col)
	plt.tight_layout()
	plt.show()
	


# In[68]:


images,labels=next(iter(dl_train))
show_images(images,labels,labels)


# In[69]:


images,labels=next(iter(dl_test))
show_images(images,labels,labels)


# # Creating the Model

# In[71]:


resnet18=torchvision.models.resnet18(pretrained=True)
print(resnet18)


# In[72]:


#To change the FC layer in Resnet18 1000 to 3 OP features

resnet18.fc=torch.nn.Linear(in_features=512,out_features=3,bias=True)
loss_fn=torch.nn.CrossEntropyLoss() #Loss Function
optimizer=torch.optim.Adam(resnet18.parameters(),lr=3e-5) # Optimization Stochastic Gradient Descent



# In[76]:


def show_preds():
	resnet18.eval()
	images,labels=next(iter(dl_test))
	outputs=resnet18(images)
	_, preds=torch.max(outputs,1)
	show_images(images,labels,preds)


# In[77]:


show_preds()


# # Training the Model

# In[78]:


def train(epochs):
	print('Starting Training...')
	for e in range(0,epochs):
		print('='*20)
		print(f'Starting epoch{e+1}/{epochs}')
		print('='*20)
		
		train_loss=0
		
		resnet18.train()
		for train_step,(images,labels) in enumerate(dl_train):
			optimizer.zero_grad()
			outputs=resnet18(images)
			loss=loss_fn(outputs,labels)
			loss.backward()
			optimizer.step()
			train_loss+=loss.item()
			if train_step%20==0:
				print('Evaluating at Step',train_step)
				acc=0.
				val_loss=0.
				resnet18.eval()
				
				for val_step,(images,labels) in enumerate(dl_test):
					outputs=resnet18(images)
					loss=loss_fn(outputs,labels)
					val_loss+=loss.item()
					
					_,preds=torch.max(outputs,1)
					acc+=sum((preds==labels).numpy())
				val_loss/=(val_step+1)
				acc=acc/len(test_dataset)
				print(f'Val Loss:{val_loss:.4f},Acc:{acc:.4f}')
				show_preds()
				
				resnet18.train()
				
				if acc>0.95:
					print('Performance Condition Satified')
					return
		train_loss/=(train_step+1)
		print(f'Training loss:{train_loss:.4f}')
		


# In[79]:


train(epochs=1)


# # Final Results

# In[80]:


show_preds()


# In[ ]:





# In[ ]:
