Update models/resnet50.py

models/resnet50.py

@@ -8,17 +8,8 @@ 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 create_indices(labels, mapping):
+    return [mapping[label] for label in labels]
 
 def write_to_csv(predicted, actual, probs, write_path, header):
 
@@ -42,7 +33,7 @@ def write_to_csv(predicted, actual, probs, write_path, header):
 
 class ResNet50():
 
-    def __init__(self, num_classes, lr=0.01, momentum=0.9):
+    def __init__(self, num_classes, lr=0.01, momentum=0.9, mapping=None):
         self.model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
         self.num_classes = num_classes
         self.lr = lr
@@ -53,6 +44,8 @@ class ResNet50():
         self.criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
         self.optimizer = optim.SGD(self.model.parameters(), lr=self.lr, momentum=self.momentum)
 
+        self.mapping = mapping
+
     def train(self, epochs, train_loader):
         loss_over_time = []
         num_epochs = list(range(1, epochs + 1))
@@ -63,7 +56,7 @@ class ResNet50():
                 inputs, labels = data
                 self.optimizer.zero_grad()
                 outputs = self.model(data[inputs].float())
-                indices = create_indices(data[labels])
+                indices = create_indices(data[labels], self.mapping)
                 target = torch.tensor(indices)
                 loss = self.criterion(outputs, target)
                 loss.backward()
@@ -79,13 +72,6 @@ class ResNet50():
             "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
@@ -97,7 +83,8 @@ class ResNet50():
                 images, labels = data
                 images = data[images].float()
                 labels = data[labels]
-                indices = create_indices(labels)
+                # indices = create_indices(labels)
+                indices = create_indices(labels, self.mapping)
                 labels = torch.tensor(indices)
 
                 outputs = self.model(images)