initial commit

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+model_4epochs_90acc.pth filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,51 @@
+import gradio as gr
+import os
+import torch
+from timeit import default_timer as timer
+from typing import Tuple, Dict
+
+from model import create_model
+
+#get the class names
+classes = ['buildings', 'forest', 'glacier', 'mountain', 'sea', 'street']
+
+#get the model and load its trained weights
+model, transform = create_model(num_classes = len(classes))
+model.load_state_dict(torch.load(f = "model_4epochs_90acc.pth" , map_location = torch.device("cpu")))
+
+#prediction function for a single image
+def predict(img):
+    start_time = timer()
+    img = transform(img).unsqueeze(0)
+    
+    #get the prediction probabilities and put them in a dictionary
+    model.eval()
+    with torch.inference_mode():
+        y_prob = model(img).softmax(dim = 1)
+        y_preds = {classes[i] : float(y_prob[0][i]) for i in range(len(classes)) }
+        
+    prediction_time = round(timer() - start_time, 5)
+    
+    return y_preds, prediction_time
+
+
+#Gradio App ###
+title = "Intel Scenery Classification"
+description = "An efficientnet_b2 model for the classification of different image scenes from an intel dataset."
+article = "Created by me, Richard Schattner."
+
+#get the examples list
+example_list = [["examples/" + example] for example in os.listdir("examples")]
+
+# Create the Gradio demo
+demo = gr.Interface(fn=predict, 
+                    inputs=gr.Image(type="pil"),
+                    outputs=[gr.Label(num_top_classes=6, label="Predictions"),
+                             gr.Number(label="Prediction time (s)")], 
+                    examples=example_list, 
+                    title=title,
+                    description=description,
+                    article=article)
+
+# Launch the demo
+demo.launch()

data.py

@@ -0,0 +1,20 @@
+from torch.utils.data import DataLoader
+import torchvision
+
+#get the correct transform for the effnet_b2 model
+weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
+transform = weights.transforms()
+
+#create test/train datasets and dataloaders
+train_dir = "intel_image/seg_train"
+test_dir = "intel_image/seg_test"
+
+train_data = torchvision.datasets.ImageFolder(root = train_dir, transform = transform)
+test_data = torchvision.datasets.ImageFolder(root = test_dir, transform = transform)
+
+train_loader = DataLoader(train_data, shuffle = True, batch_size = 32)
+test_loader = DataLoader(test_data, shuffle = False, batch_size = 32)
+
+def create_dataloaders():
+    """Returns: Training and test dataloaders """
+    return train_loader, test_loader

model.py

@@ -0,0 +1,26 @@
+import torch
+import torchvision
+from torch import nn
+
+def create_model(num_classes = 6, seed = 1):
+    """Create an instance of the effnet_b2 model, freezes all layers and changes the classifier head.
+        
+        Returns: The model and its data transform
+    """
+    #get pretrained model and its transform
+    weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
+    model = torchvision.models.efficientnet_b2(weights = weights)
+    transform = weights.transforms()
+
+    #freeze all layers
+    for param in model.parameters():
+        param.requires_grad = False
+
+    #create a new classifier head with 6 output classes
+    classifier = nn.Sequential(nn.Dropout(p = 0.2, inplace = True),
+                               nn.Linear(in_features = 1408, out_features = num_classes))
+
+    #replace old classifier head with newly created one
+    model.classifier = classifier
+    
+    return model, transform

model_4epochs_90acc.pth

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

requirements.txt

@@ -0,0 +1,3 @@
+torch == 2.3.0
+torchvision == 0.18.0
+gradio == 4.31.4

train.py

@@ -0,0 +1,92 @@
+from tqdm.auto import tqdm
+import torch
+from torch import nn
+
+from data.py import create_dataloaders
+
+#get the train/test dataloaders from data.py
+train_loader, test_loader = create_dataloaders()
+
+#define an accuracy function
+def accuracy_fn(y_true, y_pred):
+    correct = torch.eq(y_true, y_pred).sum().item()
+    acc = (correct / len(y_pred)) * 100 
+    return acc
+
+#instantiate loss function and optimizer
+loss_fn = nn.CrossEntropyLoss()
+optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3)
+
+
+#create a function for a training step
+def train_step(model):
+    train_loss, train_accuracy = 0, 0
+    model.train()
+
+    for batch, (x,y) in enumerate(train_loader):    
+        #get predictions
+        y_logits = model(x)
+        y_pred = y_logits.argmax(dim = 1)
+    
+        #calculate loss
+        loss = loss_fn(y_logits, y)
+        train_loss += loss.item()
+        train_accuracy += accuracy_fn(y, y_pred)
+        
+        #update model
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        
+    #divide test loss and accuracy by length of dataloader
+    train_loss /= len(train_loader)
+    train_accuracy /= len(train_loader)
+    
+    #return train loss and accuracy
+    return train_loss, train_accuracy
+    
+#create a function to test the model
+def test_step(model):
+    test_loss, test_accuracy = 0, 0
+    
+    model.eval()
+    with torch.inference_mode():
+        for batch, (x,y) in enumerate(test_loader):
+            y_logits = model(x)
+            y_pred = y_logits.argmax(dim = 1)
+            
+            loss = loss_fn(y_logits, y)
+            test_loss += loss.item()         
+            test_accuracy += accuracy_fn(y, y_pred)
+        
+        #divide test loss and accuracy by length of dataloader
+        test_loss /= len(test_loader)
+        test_accuracy /= len(test_loader)
+    
+    #return test loss and accuracy
+    return test_loss, test_accuracy
+
+def train(model, epochs):
+    """Trains a model for a given number of epochs
+    
+    Args: model and epochs
+    Returns: The trained model and a dictionary of train/test loss and train/test accuracy for each epoch.
+    """
+    #create an empty list of train/test metrics
+    train_loss, test_loss, train_acc, test_acc = [], [], [], []
+    for epoch in tqdm(range(epochs)):
+        #train step and save the loss and accuracy
+        new_train_loss, new_train_acc = train_step(model)
+        train_loss.append(new_train_loss)
+        train_acc.append(new_train_acc)
+        
+        #test step and save the loss and accuracy
+        new_test_loss, new_test_acc = test_step(model)
+        test_loss.append(new_test_loss)
+        test_acc.append(new_test_acc)
+        
+    #put the metrics in a dictionary
+    metrics = {"train_loss": train_loss, "test_loss" : test_loss, 
+               "train_acc": train_acc, "test_acc": test_acc}
+    
+    return model, metrics