Added preprocessing script

models/model_weights.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3cbeb72bb5c303d0cf32c7f4003990b4de68460b342d0d1cecb3e147c1f1ebd6
+size 94380387

models/natural_disaster_dataset.py

@@ -0,0 +1,107 @@
+from PIL import Image
+import os
+from torch.utils.data import DataLoader, Dataset
+import torch
+from skimage import transform
+import matplotlib.pyplot as plt
+import numpy as np
+import torchvision.transforms as transforms
+import torchvision.transforms.functional as TF
+import streamlit as st
+
+class NaturalDisasterDataset(Dataset):
+    """
+    A custom PyTorch Dataset that contains images of several types of natural disasters, 
+    including earthquakes, fires, and floods.
+    """
+    def __init__(self, root:str, transform:any=None) -> None:
+        """
+        Creates a custom PyTorch dataset of natural disasters.
+
+        Args:
+            root (str): A path containing the images.
+            transform (any): A type of transformation from the scikit-image library.
+
+        Returns:
+            None
+        """
+        self.root = root
+        self.transform = transform
+
+        self.image_paths = []
+        self.labels = []
+
+        for label in os.listdir(root):
+            folder = os.path.join(root, label)
+            for file in os.listdir(folder):
+                self.image_paths.append(os.path.join(folder, file))
+                self.labels.append(label)
+
+    def __len__(self) -> int:
+        """
+        Returns the length/size of the dataset.
+
+        Args:
+            None
+
+        Returns:
+            length (int): The length of the dataset.
+        """
+        return len(self.image_paths)
+
+    def __getitem__(self, idx:int) -> dict:
+        """
+        Iterates through the dataset and returns a sample image.
+
+        Args:
+            idx (int): An index to the dataset.
+
+        Returns:
+            sample (dict): A dictionary containing the image and its label.
+        """
+        img_path = self.image_paths[idx]
+        label = self.labels[idx]
+        image = Image.open(img_path).convert("RGB")
+
+        if self.transform:
+            image = self.transform(image) 
+
+        image = transforms.PILToTensor()(image)
+        sample = {"image": image, "category": label}
+        return sample
+
+    
+    def load_sample(self) -> None:
+        """
+        Displays four sample images, one of each type of disaster.
+
+        Args:
+            None
+
+        Returns:
+            None
+        """
+
+        categories_needed = {"Normal", "Earthquake", "Fire", "Flood"}
+        shown = {}
+
+        fig = plt.figure(figsize=(10, 3))
+
+        for sample in self:
+            category = sample["category"]
+
+            # If we still need this category
+            if category in categories_needed and category not in shown:
+                shown[category] = sample["image"]
+
+            # Stop if we have all 4 categories
+            if len(shown) == len(categories_needed):
+                break
+
+        for i, (category, image) in enumerate(shown.items()):
+            ax = plt.subplot(1, 4, i + 1)
+            ax.imshow(image)
+            ax.set_title(category)
+            ax.axis("off")
+
+        plt.tight_layout()

models/resnet50.py

@@ -0,0 +1,113 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torchvision.models as models
+import numpy as np
+from PIL import Image
+import torchvision.transforms as transforms
+import os
+import pandas as pd
+
+
+def create_indices(labels):
+    mapping = {
+        "Non_Damage": 0,
+        "Land_Disaster": 1,
+        "Fire_Disaster": 2,
+        "Water_Disaster": 3
+    }
+
+    indices = list(mapping[category] for category in labels)
+    return indices
+
+def write_to_csv(predicted, actual, probs, write_path, header):
+
+    label_names = ["Non-Damage", "Earthquake", "Fire", "Flood"]
+
+    if header:
+        with open(write_path, "w") as file:
+            file.write("Predicted,True,Non_Damage_Score,Earthquake_Score,Fire_Score,Flood_Score\n")
+
+    with open(write_path, "a") as file:
+        for i in range(len(actual)):
+            file.write(
+                f"{label_names[actual[i].item()]},"
+                f"{label_names[predicted[i].item()]},"
+                f"{probs[i, 0].item()},"
+                f"{probs[i, 1].item()},"
+                f"{probs[i, 2].item()},"
+                f"{probs[i, 3].item()}\n"
+            )
+
+
+class ResNet50():
+
+    def __init__(self, num_classes, lr=0.01, momentum=0.9):
+        self.model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
+        self.num_classes = num_classes
+        self.lr = lr
+        self.momentum = momentum
+        self.num_features = self.model.fc.in_features
+        self.model.fc = nn.Linear(self.num_features, self.num_classes)
+
+        self.criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
+        self.optimizer = optim.SGD(self.model.parameters(), lr=self.lr, momentum=self.momentum)
+
+    def train(self, epochs, train_loader):
+        loss_over_time = []
+        num_epochs = list(range(1, epochs + 1))
+        for epoch in range(epochs):
+            self.model.train()
+            current_loss = 0.0
+            for i, data in enumerate(train_loader, 0):
+                inputs, labels = data
+                self.optimizer.zero_grad()
+                outputs = self.model(data[inputs].float())
+                indices = create_indices(data[labels])
+                target = torch.tensor(indices)
+                loss = self.criterion(outputs, target)
+                loss.backward()
+                self.optimizer.step()
+                current_loss += loss.item()
+            loss_over_time.append(current_loss / len(train_loader))
+            print(f"Epoch: {epoch + 1} \t Loss: {current_loss / len(train_loader)}")
+
+        torch.save({
+            "model_state_dict": self.model.state_dict(),
+            "optimizer_state_dict": self.optimizer.state_dict(),
+            "epochs": num_epochs,
+            "loss": loss_over_time
+        }, "model_weights.pth")
+
+        data = {
+            "Epochs": num_epochs,
+            "Loss": loss_over_time
+        }
+        data = pd.DataFrame(data=data)
+        data.to_csv("results/model_progress.csv", index=False)
+
+    def eval(self, test_loader, write_path=None):
+        self.model.eval()
+        header = True
+
+        with torch.no_grad():
+            correct = 0
+            total = 0
+            for data in test_loader:
+                images, labels = data
+                images = data[images].float()
+                labels = data[labels]
+                indices = create_indices(labels)
+                labels = torch.tensor(indices)
+
+                outputs = self.model(images)
+                _, predicted = torch.max(outputs.data, 1)
+                probs = torch.softmax(outputs, dim=1)
+
+                total += len(labels)
+                correct += (predicted == labels).sum().item()
+                if write_path:
+                    write_to_csv(predicted, labels, probs, write_path=write_path, header=header)
+                header = False
+        
+        print(f'Accuracy of the network on the test images: {round(100 * correct / total, 3)}%')

models/save_model_progress.py

@@ -0,0 +1,19 @@
+import pandas as pd
+from resnet50 import ResNet50
+import torch
+import matplotlib.pyplot as plt
+
+n_categories = 4
+model = ResNet50(n_categories)
+weights = torch.load("model_weights.pth")
+model.model.load_state_dict(weights["model_state_dict"])
+epochs = weights["epochs"]
+loss = weights["loss"]
+
+data = {
+    "Epochs": epochs,
+    "Loss": loss
+}
+
+data = pd.DataFrame(data)
+data.to_csv("model_progress.csv", index=False)

models/upload_model.py

@@ -0,0 +1,15 @@
+import torch
+from resnet50 import ResNet50
+from huggingface_hub import HfApi
+
+api = HfApi()
+
+model = ResNet50(num_classes=4)
+torch.save(model.model.state_dict(), "model_weights.pth")
+
+api.upload_file(
+    path_or_fileobj="model_weights.pth",
+    path_in_repo="model_weights.pth",
+    repo_id="DanielCruz09/disaster-image-classifier",
+    repo_type="model"
+)