Update

main.py

@@ -0,0 +1,217 @@
+import os
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision.transforms import transforms
+from torch.utils.data import DataLoader, random_split, Dataset
+from torchvision.datasets import ImageFolder
+import matplotlib.pyplot as plt
+from models import *
+from scipy.ndimage import gaussian_filter1d
+import numpy as np
+
+# Constants
+RANDOM_SEED = 123
+BATCH_SIZE = 32
+NUM_EPOCHS = 100
+LEARNING_RATE = 0.0001
+STEP_SIZE = 10
+GAMMA = 0.5
+DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+NUM_PRINT = 100
+NUM_CLASSES = 5
+
+# Load and preprocess the data
+data_dir = r"data/train/Task 1"
+
+# Define transformation for preprocessing
+preprocess = transforms.Compose(
+    [
+        transforms.Resize((64, 64)),  # Resize images to 64x64
+        transforms.ToTensor(),  # Convert to tensor
+        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),  # Normalize
+    ]
+)
+
+augmentation = transforms.Compose(
+    [
+        transforms.Resize((64, 64)),  # Resize images to 64x64
+        transforms.RandomHorizontalFlip(p=0.5),  # Random horizontal flip
+        transforms.RandomRotation(degrees=45),  # Random rotation
+        transforms.RandomVerticalFlip(p=0.5),  # Random vertical flip
+        transforms.RandomGrayscale(p=0.1),  # Random grayscale
+        transforms.ColorJitter(
+            brightness=0.5, contrast=0.5, saturation=0.5, hue=0.5
+        ),  # Random color jitter
+        transforms.ToTensor(),  # Convert to tensor
+        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),  # Normalize
+    ]
+)
+
+# Load the dataset using ImageFolder
+original_dataset = ImageFolder(root=data_dir, transform=preprocess)
+augmented_dataset = ImageFolder(root=data_dir, transform=augmentation)
+dataset = original_dataset + augmented_dataset
+
+print("Length of dataset: ", len(dataset))
+print("Classes: ", original_dataset.classes)
+
+
+# Custom dataset class
+class CustomDataset(Dataset):
+    def __init__(self, dataset):
+        self.data = dataset
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        img, label = self.data[idx]
+        return img, label
+
+
+# Split the dataset into train and validation sets
+train_size = int(0.8 * len(dataset))
+val_size = len(dataset) - train_size
+train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
+
+# Create data loaders for the custom dataset
+train_loader = DataLoader(
+    CustomDataset(train_dataset), batch_size=BATCH_SIZE, shuffle=True, num_workers=0
+)
+valid_loader = DataLoader(
+    CustomDataset(val_dataset), batch_size=BATCH_SIZE, num_workers=0
+)
+
+# Initialize model, criterion, optimizer, and scheduler
+model = resnet18(pretrained=False, num_classes=NUM_CLASSES)
+model = model.to(DEVICE)
+criterion = nn.CrossEntropyLoss()
+# Adam optimizer
+optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
+# ReduceLROnPlateau scheduler
+scheduler = optim.lr_scheduler.ReduceLROnPlateau(
+    optimizer, mode="min", factor=0.1, patience=10, verbose=True
+)
+
+# Lists to store training and validation loss history
+TRAIN_LOSS_HIST = []
+VAL_LOSS_HIST = []
+AVG_TRAIN_LOSS_HIST = []
+AVG_VAL_LOSS_HIST = []
+TRAIN_ACC_HIST = []
+VAL_ACC_HIST = []
+
+# Training loop
+for epoch in range(NUM_EPOCHS):
+    model.train(True)  # Set model to training mode
+    running_loss = 0.0
+    total_train = 0
+    correct_train = 0
+
+    for i, (inputs, labels) in enumerate(train_loader, 0):
+        inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
+        optimizer.zero_grad()
+        outputs = model(inputs)
+        loss = criterion(outputs, labels)
+        loss.backward()
+        optimizer.step()
+        running_loss += loss.item()
+
+        if (i + 1) % NUM_PRINT == 0:
+            print(
+                "[Epoch %d, Batch %d] Loss: %.6f"
+                % (epoch + 1, i + 1, running_loss / NUM_PRINT)
+            )
+            running_loss = 0.0
+
+        _, predicted = torch.max(outputs, 1)
+        total_train += labels.size(0)
+        correct_train += (predicted == labels).sum().item()
+
+    TRAIN_LOSS_HIST.append(loss.item())
+
+    # Calculate the average training loss for the epoch
+    avg_train_loss = running_loss / len(train_loader)
+    AVG_TRAIN_LOSS_HIST.append(avg_train_loss)
+
+    # Print average training loss for the epoch
+    print("[Epoch %d] Average Training Loss: %.6f" % (epoch + 1, avg_train_loss))
+
+    # Learning rate scheduling
+    lr_1 = optimizer.param_groups[0]["lr"]
+    print("Learning Rate: {:.15f}".format(lr_1))
+    scheduler.step(avg_val_loss)
+
+    # Validation loop
+    model.eval()  # Set model to evaluation mode
+    val_loss = 0.0
+    correct_val = 0
+    total_val = 0
+
+    with torch.no_grad():
+        for inputs, labels in valid_loader:
+            inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
+            outputs = model(inputs)
+            loss = criterion(outputs, labels)
+            val_loss += loss.item()
+            # Calculate accuracy
+            _, predicted = torch.max(outputs, 1)
+            total_val += labels.size(0)
+            correct_val += (predicted == labels).sum().item()
+
+    VAL_LOSS_HIST.append(loss.item())
+
+    # Calculate the average validation loss for the epoch
+    avg_val_loss = val_loss / len(valid_loader)
+    AVG_VAL_LOSS_HIST.append(loss.item())
+    print("Average Validation Loss: %.6f" % (avg_val_loss))
+
+    # Calculate the accuracy of validation set
+    val_accuracy = correct_val / total_val
+    VAL_ACC_HIST.append(val_accuracy)
+    print("Validation Accuracy: %.6f" % (val_accuracy))
+
+# End of training loop
+
+# Save the model
+model_save_path = "model.pth"
+torch.save(model.state_dict(), model_save_path)
+print("Model saved at", model_save_path)
+
+print("Generating loss plot...")
+# Make the plot smoother by interpolating the data
+# https://stackoverflow.com/questions/5283649/plot-smooth-line-with-pyplot
+# train_loss_line = gaussian_filter1d(TRAIN_LOSS_HIST, sigma=10)
+# val_loss_line = gaussian_filter1d(VAL_LOSS_HIST, sigma=10)
+# plt.plot(range(1, NUM_EPOCHS + 1), train_loss_line, label='Train Loss')
+# plt.plot(range(1, NUM_EPOCHS + 1), val_loss_line, label='Validation Loss')
+avg_train_loss_line = gaussian_filter1d(AVG_TRAIN_LOSS_HIST, sigma=2)
+avg_val_loss_line = gaussian_filter1d(AVG_VAL_LOSS_HIST, sigma=2)
+train_loss_line = gaussian_filter1d(TRAIN_LOSS_HIST, sigma=2)
+val_loss_line = gaussian_filter1d(VAL_LOSS_HIST, sigma=2)
+train_acc_line = gaussian_filter1d(TRAIN_ACC_HIST, sigma=2)
+val_acc_line = gaussian_filter1d(VAL_ACC_HIST, sigma=2)
+plt.plot(range(1, NUM_EPOCHS + 1), train_loss_line, label="Train Loss")
+plt.plot(range(1, NUM_EPOCHS + 1), val_loss_line, label="Validation Loss")
+plt.xlabel("Epochs")
+plt.ylabel("Loss")
+plt.legend()
+plt.title("Train Loss and Validation Loss")
+plt.savefig("loss_plot.png")
+plt.clf()
+plt.plot(range(1, NUM_EPOCHS + 1), avg_train_loss_line, label="Average Train Loss")
+plt.plot(range(1, NUM_EPOCHS + 1), avg_val_loss_line, label="Average Validation Loss")
+plt.xlabel("Epochs")
+plt.ylabel("Loss")
+plt.legend()
+plt.title("Average Train Loss and Average Validation Loss")
+plt.savefig("avg_loss_plot.png")
+plt.clf()
+plt.plot(range(1, NUM_EPOCHS + 1), train_acc_line, label="Train Accuracy")
+plt.plot(range(1, NUM_EPOCHS + 1), val_acc_line, label="Validation Accuracy")
+plt.xlabel("Epochs")
+plt.ylabel("Accuracy")
+plt.legend()
+plt.title("Train Accuracy and Validation Accuracy")
+plt.savefig("accuracy_plot.png")

models.py

@@ -0,0 +1,340 @@
+#######################################################
+# This file stores all the models used in the project.#
+#######################################################
+
+import torch
+from torchvision.models import resnet50
+from torchvision.models import resnet18
+
+# resnet50
+class Bottleneck(torch.nn.Module):
+    expansion = 4
+
+    def __init__(self, in_channels, out_channels, i_downsample=None, stride=1):
+        super(Bottleneck, self).__init__()
+        # hmm,ex 1x1 convolution to reduce channels (intermediate channels)
+        self.conv1 = torch.nn.Conv2d(
+            in_channels, out_channels, kernel_size=1, stride=1, padding=0
+        )
+        self.batch_norm1 = torch.nn.BatchNorm2d(out_channels)
+        # 3x3 convolution with specified stride
+        self.conv2 = torch.nn.Conv2d(
+            out_channels, out_channels, kernel_size=3, stride=stride, padding=1
+        )
+        self.batch_norm2 = torch.nn.BatchNorm2d(out_channels)
+        # and then leh,1x1 expand back
+        self.conv3 = torch.nn.Conv2d(
+            out_channels,
+            out_channels * self.expansion,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        )
+        self.batch_norm3 = torch.nn.BatchNorm2d(out_channels * self.expansion)
+
+        self.i_downsample = i_downsample
+        self.stride = stride
+        self.relu = torch.nn.ReLU()
+
+    ##forward the input x through the network,haiyaa
+    def forward(self, x):
+        identity = x.clone()
+        x = self.relu(self.batch_norm1(self.conv1(x)))
+
+        x = self.relu(self.batch_norm2(self.conv2(x)))
+
+        x = self.conv3(x)
+        x = self.batch_norm3(x)
+
+        # downsample if needed
+        if self.i_downsample is not None:
+            identity = self.i_downsample(identity)
+        # add identity
+        x += identity
+        x = self.relu(x)
+
+        return x
+
+
+# we no use this first,but we can just copy this whole class and apply to resnet16 and etc
+class Block(torch.nn.Module):
+    expansion = 1
+
+    def __init__(self, in_channels, out_channels, i_downsample=None, stride=1):
+        super(Block, self).__init__()
+
+        self.conv1 = torch.nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size=3,
+            padding=1,
+            stride=stride,
+            bias=False,
+        )
+        self.batch_norm1 = torch.nn.BatchNorm2d(out_channels)
+        self.conv2 = torch.nn.Conv2d(
+            out_channels,
+            out_channels,
+            kernel_size=3,
+            padding=1,
+            stride=stride,
+            bias=False,
+        )
+        self.batch_norm2 = torch.nn.BatchNorm2d(out_channels)
+
+        self.i_downsample = i_downsample
+        self.stride = stride
+        self.relu = torch.nn.ReLU()
+
+    def forward(self, x):
+        identity = x.clone()
+
+        x = self.relu(self.batch_norm2(self.conv1(x)))
+        x = self.batch_norm2(self.conv2(x))
+
+        if self.i_downsample is not None:
+            identity = self.i_downsample(identity)
+        print(x.shape)
+        print(identity.shape)
+        x += identity
+        x = self.relu(x)
+        return x
+
+
+class ResNet(torch.nn.Module):
+    def __init__(self, ResBlock, layer_list, num_classes, num_channels=3):
+        super(ResNet, self).__init__()
+        self.in_channels = 64
+        # intial conv layaer
+        self.conv1 = torch.nn.Conv2d(
+            num_channels, 64, kernel_size=7, stride=2, padding=3, bias=False
+        )
+        self.batch_norm1 = torch.nn.BatchNorm2d(64)
+        self.relu = torch.nn.ReLU()
+        self.max_pool = torch.nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        # residual block(layers),each block got three three layer,total 4 blocks
+        self.layer1 = self._make_layer(ResBlock, layer_list[0], planes=64)
+        self.layer2 = self._make_layer(ResBlock, layer_list[1], planes=128, stride=2)
+        self.layer3 = self._make_layer(ResBlock, layer_list[2], planes=256, stride=2)
+        self.layer4 = self._make_layer(ResBlock, layer_list[3], planes=512, stride=2)
+
+        self.avgpool = torch.nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = torch.nn.Linear(512 * ResBlock.expansion, num_classes)
+
+    def forward(self, x):
+        x = self.relu(self.batch_norm1(self.conv1(x)))
+        x = self.max_pool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = x.reshape(x.shape[0], -1)
+        x = self.fc(x)
+
+        return x
+
+    def _make_layer(self, ResBlock, blocks, planes, stride=1):
+        # plane is the number of output channel
+        ii_downsample = None
+        layers = []
+
+        if stride != 1 or self.in_channels != planes * ResBlock.expansion:
+            ii_downsample = torch.nn.Sequential(
+                torch.nn.Conv2d(
+                    self.in_channels,
+                    planes * ResBlock.expansion,
+                    kernel_size=1,
+                    stride=stride,
+                ),
+                torch.nn.BatchNorm2d(planes * ResBlock.expansion),
+            )
+
+        layers.append(
+            ResBlock(
+                self.in_channels, planes, i_downsample=ii_downsample, stride=stride
+            )
+        )
+        self.in_channels = planes * ResBlock.expansion
+
+        for i in range(blocks - 1):
+            layers.append(ResBlock(self.in_channels, planes))
+
+        return torch.nn.Sequential(*layers)
+
+
+##list here leh is the number of residual block in each layer
+def ResNet50(num_classes, channels=3):
+    return ResNet(Bottleneck, [3, 4, 6, 3], num_classes, channels)
+
+
+# VGG16 model
+class VGG16(torch.nn.Module):
+    def __init__(self, num_classes):
+        super().__init__()
+
+        self.block_1 = torch.nn.Sequential(
+            torch.nn.Conv2d(
+                in_channels=3,
+                out_channels=64,
+                kernel_size=(3, 3),
+                stride=(1, 1),
+                padding=1,
+            ),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(
+                in_channels=64,
+                out_channels=64,
+                kernel_size=(3, 3),
+                stride=(1, 1),
+                padding=1,
+            ),
+            torch.nn.ReLU(),
+            torch.nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2)),
+        )
+
+        self.block_2 = torch.nn.Sequential(
+            torch.nn.Conv2d(
+                in_channels=64,
+                out_channels=128,
+                kernel_size=(3, 3),
+                stride=(1, 1),
+                padding=1,
+            ),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(
+                in_channels=128,
+                out_channels=128,
+                kernel_size=(3, 3),
+                stride=(1, 1),
+                padding=1,
+            ),
+            torch.nn.ReLU(),
+            torch.nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2)),
+        )
+
+        self.block_3 = torch.nn.Sequential(
+            torch.nn.Conv2d(
+                in_channels=128,
+                out_channels=256,
+                kernel_size=(3, 3),
+                stride=(1, 1),
+                padding=1,
+            ),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(
+                in_channels=256,
+                out_channels=256,
+                kernel_size=(3, 3),
+                stride=(1, 1),
+                padding=1,
+            ),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(
+                in_channels=256,
+                out_channels=256,
+                kernel_size=(3, 3),
+                stride=(1, 1),
+                padding=1,
+            ),
+            torch.nn.ReLU(),
+            torch.nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2)),
+        )
+
+        self.block_4 = torch.nn.Sequential(
+            torch.nn.Conv2d(
+                in_channels=256,
+                out_channels=512,
+                kernel_size=(3, 3),
+                stride=(1, 1),
+                padding=1,
+            ),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(
+                in_channels=512,
+                out_channels=512,
+                kernel_size=(3, 3),
+                stride=(1, 1),
+                padding=1,
+            ),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(
+                in_channels=512,
+                out_channels=512,
+                kernel_size=(3, 3),
+                stride=(1, 1),
+                padding=1,
+            ),
+            torch.nn.ReLU(),
+            torch.nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2)),
+        )
+
+        self.block_5 = torch.nn.Sequential(
+            torch.nn.Conv2d(
+                in_channels=512,
+                out_channels=512,
+                kernel_size=(3, 3),
+                stride=(1, 1),
+                padding=1,
+            ),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(
+                in_channels=512,
+                out_channels=512,
+                kernel_size=(3, 3),
+                stride=(1, 1),
+                padding=1,
+            ),
+            torch.nn.ReLU(),
+            torch.nn.Conv2d(
+                in_channels=512,
+                out_channels=512,
+                kernel_size=(3, 3),
+                stride=(1, 1),
+                padding=1,
+            ),
+            torch.nn.ReLU(),
+            torch.nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2)),
+        )
+
+        height, width = 3, 3
+        self.classifier = torch.nn.Sequential(
+            torch.nn.Linear(512 * height * width, 4096),
+            torch.nn.ReLU(True),
+            torch.nn.Dropout(p=0.5),
+            torch.nn.Linear(4096, 4096),
+            torch.nn.ReLU(True),
+            torch.nn.Dropout(p=0.5),
+            torch.nn.Linear(4096, num_classes),
+        )
+
+        for m in self.modules():
+            if isinstance(m, torch.torch.nn.Conv2d) or isinstance(
+                m, torch.torch.nn.Linear
+            ):
+                torch.nn.init.kaiming_uniform_(
+                    m.weight, mode="fan_in", nonlinearity="relu"
+                )
+                if m.bias is not None:
+                    m.bias.detach().zero_()
+
+        self.avgpool = torch.nn.AdaptiveAvgPool2d((height, width))
+
+    def forward(self, x):
+        x = self.block_1(x)
+        x = self.block_2(x)
+        x = self.block_3(x)
+        x = self.block_4(x)
+        x = self.block_5(x)
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)  # flatten
+
+        logits = self.classifier(x)
+        # probas = F.softmax(logits, dim=1)
+
+        return logits
+
+
+# ResNet18 model

predict.py

@@ -0,0 +1,101 @@
+import os
+import torch
+import torch.nn as nn
+from torchvision import transforms
+from PIL import Image
+from models import *  # Make sure you import your model correctly from the 'models' module
+from torchmetrics import ConfusionMatrix
+import matplotlib.pyplot as plt
+import pathlib 
+
+# Define the path to your model checkpoint
+model_checkpoint_path = "model.pth"
+
+# Define the path to the image you want to classify
+image_path = "data/test/Task 1/"  # Use forward slashes for file paths
+
+# Define images variable to recursively list all the data file in the image_path
+images = list(pathlib.Path(image_path).rglob("*.png"))
+classes = os.listdir(image_path)
+print(images)
+
+true_classs = []
+predicted_labels = []
+
+DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+NUM_CLASSES = 5  # Update with the correct number of classes
+
+# Load your model (change this according to your model definition)
+model = resnet18(pretrained=False, num_classes=NUM_CLASSES)
+model.load_state_dict(torch.load(model_checkpoint_path, map_location=DEVICE))  # Load the model on the same device
+model.eval()
+model = model.to(DEVICE)
+
+# Define transformation for preprocessing the input image
+preprocess = transforms.Compose(
+    [
+        transforms.Resize((64, 64)),  # Resize the image to match training input size
+        transforms.Grayscale(num_output_channels=3),  # Convert the image to grayscale
+        transforms.ToTensor(),
+        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),  # Normalize the image
+    ]
+)
+
+def predict_image(image_path, model, transform):
+    model.eval()
+    correct_predictions = 0
+    total_predictions = len(images)
+    
+    with torch.no_grad():
+        for i in images:
+            print('---------------------------')
+            # Check the true label of the image by checking the sequence of the folder in Task 1
+            true_class = classes.index(i.parts[-2])
+            print("Image path:", i)
+            print("True class:", true_class)
+            image = Image.open(i)
+            image = transform(image).unsqueeze(0)
+            image = image.to(DEVICE)
+            output = model(image)
+
+            # softmax algorithm
+            probabilities = torch.softmax(output, dim=1)[0] * 100
+            predicted_class = torch.argmax(output, dim=1).item()
+
+            # Append true and predicted labels to their respective lists
+            true_classs.append(true_class)
+            predicted_labels.append(predicted_class)
+            
+            # Check if the prediction is correct
+            if predicted_class == true_class:
+                correct_predictions += 1
+
+            # Report the prediction
+            print("Predicted class:", predicted_class)
+            print("Probability:", probabilities[predicted_class].item())
+            print("Predicted label:", classes[predicted_class])
+            print("Correct predictions:", correct_predictions)
+            print("Correct?", "Yes" if predicted_class == true_class else "No")
+            print("---------------------------")
+            
+
+
+    # Calculate accuracy
+    accuracy = correct_predictions / total_predictions
+    print("Accuracy:", accuracy)
+
+# Call the predict_image function
+predict_image(image_path, model, preprocess)
+
+# Convert the lists to tensors
+predicted_labels_tensor = torch.tensor(predicted_labels)
+true_classs_tensor = torch.tensor(true_classs)
+
+# Create confusion matrix
+conf_matrix = ConfusionMatrix(num_classes=NUM_CLASSES, task='multiclass')
+conf_matrix.update(predicted_labels_tensor, true_classs_tensor)
+
+# Plot confusion matrix
+conf_matrix.plot()
+plt.show()