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

# In[3]:


#import modules
import torch
from torch import nn
from torch.utils.data import DataLoader, random_split, Subset
from torchvision import datasets, transforms, models
from torchvision.datasets import ImageFolder
from torchinfo import summary

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from PIL import Image

import os
from pathlib import Path
from tqdm.auto import tqdm
import random
import os
import random
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw, ImageFont

#set random seed for reproducibility
seed = 42


# In[4]:


# Setup device-agnostic code
device = "cuda" if torch.cuda.is_available() else "cpu"
device


# In[10]:


#define function that walks through each directory and subdirectory
def walk_through_dir(dir_path):
  for dirpath, dirnames, filenames in os.walk(dir_path):
    print(f"There are {len(dirnames)} directories and {len(filenames)} images in '{dirpath}'.")

walk_through_dir("/kaggle/input/fire-dataset/fire_dataset")


# In[11]:


#setup path for our data files
fire_dir = Path("/kaggle/input/fire-dataset/fire_dataset/fire_images")
non_fire_dir = Path("/kaggle/input/fire-dataset/fire_dataset/non_fire_images")

fire_dir, non_fire_dir


# In[12]:


#set seed
random.seed(seed)

#get all image paths from the fire file (* means "any combination")
fire_image_list = list(fire_dir.glob("*.*.png"))

#get random image path
random_fire_image_path = random.choice(fire_image_list)

#get image class from path name
fire_image_class = random_fire_image_path.parent.stem

#open image
img = Image.open(random_fire_image_path)

print(f"Random image path: {random_fire_image_path}")
print(f"Image class: {fire_image_class}")
print(f"Image height: {img.height}") 
print(f"Image width: {img.width}")
img


# In[13]:


random.seed(seed)

#get all image paths from the non-fire file (* means "any combination")
non_fire_image_list = list(non_fire_dir.glob("*.*.png"))

random_non_fire_image_path = random.choice(non_fire_image_list)
non_fire_image_class = random_non_fire_image_path.parent.stem
img = Image.open(random_non_fire_image_path)

print(f"Random image path: {random_non_fire_image_path}")
print(f"Image class: {non_fire_image_class}")
print(f"Image height: {img.height}") 
print(f"Image width: {img.width}")
img


# In[14]:


#transforms.Compose allows us to package all of the transformations we want to perform on the image
data_transform = transforms.Compose([
    #resize the images to 128x128
    transforms.Resize(size=(128, 128)),
    #randomly pick a transformation with a random magnitude and apply one transformation to one image
    transforms.TrivialAugmentWide(num_magnitude_bins=31),
    #turn the image into a torch.Tensor and also converts all pixel values from 0 to 255 to be between 0.0 and 1.0 
    transforms.ToTensor() 
])


# In[15]:


with Image.open(random_fire_image_path) as f:
    
    #before transformation
    fig, ax = plt.subplots(1, 2)
    ax[0].imshow(f) 
    ax[0].set_title(f"Original \nSize: {f.size}")
    ax[0].axis("off")
    
    #after transformation
    #PyTorch default is [C, H, W] but Matplotlib is [H, W, C]
    torch.manual_seed(seed)
    transformed_image = data_transform(f).permute(1, 2, 0) 
    ax[1].imshow(transformed_image) 
    ax[1].set_title(f"Transformed \nSize: {transformed_image.shape}")
    ax[1].axis("off")


# In[16]:


with Image.open(random_non_fire_image_path) as f:
    fig, ax = plt.subplots(1, 2)
    ax[0].imshow(f) 
    ax[0].set_title(f"Original \nSize: {f.size}")
    ax[0].axis("off")

    torch.manual_seed(seed)
    transformed_image = data_transform(f).permute(1, 2, 0) 
    ax[1].imshow(transformed_image) 
    ax[1].set_title(f"Transformed \nSize: {transformed_image.shape}")
    ax[1].axis("off")


# In[17]:


#define data root directory
root_dir = "/kaggle/input/fire-dataset/fire_dataset"

#you can also apply data transformation using the `transforms` parameter here
fire_dataset = ImageFolder(root_dir)


# In[18]:


#the classes are the names of the files in the root folder
class_names = fire_dataset.classes
class_names


# In[19]:


#class names as a dict
class_dict = fire_dataset.class_to_idx
class_dict


# In[20]:


#split the dataset into training, validation and testing
train_dataset, valid_dataset, test_dataset = random_split(fire_dataset, [0.7, 0.15, 0.15], generator=torch.Generator().manual_seed(seed))


# In[21]:


len(train_dataset), len(valid_dataset), len(test_dataset)


# In[22]:


#define custom data subset
class MySubset(torch.utils.data.Dataset):
    def __init__(self, subset, transform=None):
        self.subset = subset
        self.transform = transform
        
    def __getitem__(self, index):
        x, y = self.subset[index]
        
        #if a transform argument is provided, apply the transformaiton to the subset
        if self.transform:
            x = self.transform(x)
        return x, y
        
    def __len__(self):
        return len(self.subset)


# In[23]:


#apply data transformation only to the training and validation dataset
train_dataset = MySubset(train_dataset, transform=data_transform)
valid_dataset = MySubset(valid_dataset, transform=data_transform)


# In[24]:


#use DataLoader to load and group dataset into batches
batch_size = 32

torch.manual_seed(seed)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

torch.manual_seed(seed)
valid_loader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=False)


# In[25]:


#check that each batch contains 32 samples of images
img, label = next(iter(train_loader))

print(f"Image shape: {img.shape} -> [batch_size, color_channels, height, width]")
print(f"Label shape: {label.shape}")


# In[26]:


#define architecture for our baseline CNN model
class Baseline(nn.Module):

    def __init__(self, input_shape: int, hidden_units: int, output_shape: int) -> None:
        super().__init__()
        
        self.conv_block_1 = nn.Sequential(
            #first convolution layer
            nn.Conv2d(in_channels=input_shape, out_channels=hidden_units, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            #second convolution layer
            nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            #pooling layer
            nn.MaxPool2d(kernel_size=3, stride=3)
        )

        self.classifier = nn.Sequential(
            nn.Flatten(),
            #we can find the number of layers checking the summary from torchinfo (example below)
            nn.Linear(in_features=17640,
                      out_features=output_shape)
        )
    
    def forward(self, x: torch.Tensor):
        
        return self.classifier(self.conv_block_1(x)) #leverage the benefits of operator fusion


# In[27]:


torch.manual_seed(seed)
#instantiate baseline model
baseline = Baseline(input_shape=3,
                  hidden_units=10, 
                  output_shape=1).to(device)
baseline


# In[28]:


#see the outputs by feeding in a test pass
summary(baseline, input_size=[32, 3, 128, 128])


# In[29]:


#define function for training and evaluating each epoch
def train_per_epoch(model, train_dataloader, test_dataloader, loss_fn, optimizer):

    model.train() 
    train_loss, train_acc = 0, 0
    
    #loop through data loader data batches

    for batch, (X, y) in enumerate(train_dataloader):
        
        X, y = X.to(device), torch.unsqueeze(y, dim=1).float().to(device)
        
        #make prediction
        y_logits = model(X)
        y_pred = torch.round(torch.sigmoid(y_logits))
        
        #calculate and accumulate loss
        loss = loss_fn(y_logits, y)
        train_loss += loss.item()
        
        #calculate and accumulate accuracy
        acc = (y_pred == y).sum().item()/len(y)
        train_acc += acc
        
        #update parameters
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    #Calculate the average loss and average acc by batch (in one epoch)
    train_loss = train_loss / len(train_dataloader)
    train_acc = train_acc / len(train_dataloader)

    ### Testing Loop
    model.eval() 
    test_loss, test_acc = 0, 0
    
    with torch.inference_mode():

        for batch, (X, y) in enumerate(test_dataloader):

            X, y = X.to(device), torch.unsqueeze(y, dim=1).float().to(device)

            test_logits = model(X)
            test_pred = torch.round(torch.sigmoid(test_logits))

            loss = loss_fn(test_logits, y)
            test_loss += loss.item()

            #calculate and accumulate accuracy
            acc = (test_pred == y).sum().item()/len(y)
            test_acc += acc
            
    #Calculate the average loss and average acc by batch (in one epoch)
    test_loss = test_loss / len(test_dataloader)
    test_acc = test_acc / len(test_dataloader)
    return train_loss, train_acc, test_loss, test_acc


# In[30]:


#setup ModelCheckpoint
def checkpoint(model, filename):
    torch.save(model.state_dict(), filename)
    
def resume(model, filename):
    model.load_state_dict(torch.load(filename))


# In[31]:


#define train function
def train(model, train_dataloader, test_dataloader, optimizer, loss_fn, epochs, best_model_path):
    
    results = {"train_loss": [], "train_acc": [], "test_loss": [], "test_acc": []}
    
    #setup early stopping and model checkpoint
    early_stop_threshold = 10
    best_loss = float("inf")
    best_epoch = -1
    
    for epoch in tqdm(range(epochs)):
        train_loss, train_acc, test_loss, test_acc = train_per_epoch(model=model, train_dataloader=train_dataloader, test_dataloader=test_dataloader, loss_fn=loss_fn, optimizer=optimizer)
        
        print(f"Epoch: {epoch+1} | train_loss: {train_loss:.4f} | train_acc: {train_acc:.4f} | test_loss: {test_loss:.4f} | test_acc: {test_acc:.4f}")

        #update results dictionary
        results["train_loss"].append(train_loss)
        results["train_acc"].append(train_acc)
        results["test_loss"].append(test_loss)
        results["test_acc"].append(test_acc)
        
        #EarlyStopping and ModelCheckpoint
        if test_loss < best_loss:
            best_loss = test_loss
            best_epoch = epoch
            checkpoint(model, best_model_path)
        elif epoch - best_epoch > early_stop_threshold:
            print("Early stopped training at epoch %d" % epoch+1)
            break  # terminate the training loop
    
    print(f"Epoch with best loss: Epoch {best_epoch+1} at {best_loss:.4f}")

    return results


# In[32]:


#set random seeds
torch.manual_seed(42) 
torch.cuda.manual_seed(42)

NUM_EPOCHS = 10

#loss function and optimizer
loss_fn = nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(params=baseline.parameters(), lr=0.001)

#train baseline model
baseline_results = train(model=baseline, train_dataloader=train_loader, test_dataloader=valid_loader, optimizer=optimizer, loss_fn=loss_fn, epochs=NUM_EPOCHS, best_model_path = "best_baseline_model.pth")


# In[33]:


def plot_results(results):
    
    train_loss = results["train_loss"]
    train_acc = results["train_acc"]
    test_loss = results["test_loss"]
    test_acc = results["test_acc"]
    
    epochs = range(len(results['train_loss']))
    
    plt.figure(figsize=(10, 5))
    
    #plot loss
    plt.subplot(1, 2, 1)
    plt.plot(epochs, train_loss, label='train_loss')
    plt.plot(epochs, test_loss, label='test_loss')
    plt.title('Loss')
    plt.xlabel('Epochs')
    plt.legend()

    #plot accuracy
    plt.subplot(1, 2, 2)
    plt.plot(epochs, train_acc, label='train_accuracy')
    plt.plot(epochs, test_acc, label='test_accuracy')
    plt.title('Accuracy')
    plt.xlabel('Epochs')
    plt.legend();


# In[34]:


plot_results(baseline_results)


# In[35]:


#get a set of pretrained model weights
weights = models.MobileNet_V3_Small_Weights.DEFAULT
weights


# In[36]:


#get the data transformations used to create our pretrained weights
auto_transforms = weights.transforms()
auto_transforms


# In[37]:


root_dir = "/kaggle/input/fire-dataset/fire_dataset"

fire_dataset_tl = ImageFolder(root_dir)


# In[38]:


#random_split the dataset in training, validation and testing
train_dataset_tl, valid_dataset_tl, test_dataset_tl = random_split(fire_dataset_tl, [0.7, 0.15, 0.15], generator=torch.Generator().manual_seed(seed))


# In[39]:


#apply data transformation to the training and validation dataset
train_dataset_tl = MySubset(train_dataset_tl, transform=auto_transforms)
valid_dataset_tl = MySubset(valid_dataset_tl, transform=auto_transforms)


# In[40]:


#load the data using DataLoader
batch_size = 32
train_loader_tl = DataLoader(train_dataset_tl, batch_size=batch_size, shuffle=True)
valid_loader_tl = DataLoader(valid_dataset_tl, batch_size=batch_size, shuffle=False)


# In[41]:


#load the weights into the model
model_tl = models.mobilenet_v3_small(weights=weights).to(device)
model_tl


# In[42]:


summary(model=model_tl, input_size=(32, 3, 384, 384))


# In[43]:


#freeze all base layers in the pretrained model
for param in model_tl.features.parameters():
    param.requires_grad = False


# In[44]:


#set the manual seeds
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)

#redefine classifier layer
model_tl.classifier = torch.nn.Sequential(
    nn.Linear(in_features=576, out_features=1024, bias=True),
    nn.Hardswish(),
    torch.nn.Dropout(p=0.2, inplace=True), 
    torch.nn.Linear(in_features=1024, out_features=1, bias=True)).to(device)


# In[45]:


#set random seeds
torch.manual_seed(seed) 
torch.cuda.manual_seed(seed)

NUM_EPOCHS = 50

#loss function and optimizer
loss_fn = nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(params=model_tl.parameters(), lr=0.0001)

#train model_tl
model_tl_results = train(model=model_tl, train_dataloader=train_loader_tl, test_dataloader=valid_loader_tl, optimizer=optimizer, loss_fn=loss_fn, epochs=NUM_EPOCHS, best_model_path = "best_model_tl.pth")


# In[46]:


plot_results(model_tl_results)


# In[5]:


#define function to predict custom images

def pred_and_plot_image(model, image_path, transform, device: torch.device=device):
    
    img = Image.open(image_path)

    image_transform = transform

    model.to(device)
    model.eval()
    with torch.inference_mode():

        transformed_image = image_transform(img).unsqueeze(dim=0)
        target_image_pred = model(transformed_image.to(device))

        target_image_pred_probs = torch.sigmoid(target_image_pred)
        target_image_pred_label = torch.round(target_image_pred_probs)
        
    plt.figure()
    plt.imshow(img)
    plt.title(f"Pred: {target_image_pred_label.cpu().numpy()[0][0]} | Prob: {target_image_pred_probs.cpu().numpy()[0][0]: .3f}")
    plt.axis(False);


# In[6]:


pred_and_plot_image(model, "./test_image", transform)