Delete src/resnet.py

src/resnet.py

@@ -1,273 +0,0 @@
-import torch
-from torchvision import models
-import torchvision.transforms as transforms
-from torchvision.transforms import InterpolationMode
-from torch.utils.data import Dataset, DataLoader
-import torch.nn as nn
-from sklearn.metrics import confusion_matrix
-import tqdm
-import matplotlib.pyplot as plt
-import numpy as np
-from src import utils
-
-plt.style.use("ggplot")
-plt.rcParams.update({"font.size": 14})
-plt.rcParams.update({"figure.autolayout": True})
-
-
-def main(batch_size=64, epochs=50, classes=("formal", "informal"), train: bool = True):
-    train_transform = transforms.Compose(
-        [
-            transforms.ToPILImage(),
-            transforms.RandomHorizontalFlip(p=0.5),
-            transforms.RandomVerticalFlip(p=0.5),
-            transforms.ColorJitter(
-                brightness=0.3, contrast=0.3, saturation=0.3, hue=0.1
-            ),
-            transforms.RandomRotation(
-                degrees=30, interpolation=InterpolationMode.BICUBIC
-            ),
-            transforms.RandomResizedCrop(
-                size=224, scale=(0.8, 1.0), interpolation=InterpolationMode.BICUBIC
-            ),
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-        ]
-    )
-
-    transform = transforms.Compose(
-        [
-            transforms.ToPILImage(),
-            transforms.Resize([256], interpolation=InterpolationMode.BICUBIC),
-            transforms.CenterCrop([224]),
-            transforms.ToTensor(),  # Converts the image to [0.0, 1.0] range
-            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-        ]
-    )
-
-    # Custom dataset class
-    class CustomDataset(Dataset):
-        def __init__(self, data, transform=None):
-            self.data = data
-            self.transform = transform
-
-        def __len__(self):
-            return len(self.data)
-
-        def __getitem__(self, idx):
-            image, label = self.data[idx]
-            if self.transform:
-                image = self.transform(image)
-            return image, label
-
-    # Load the data
-    train_data = utils.load_from_pickle("./data/high_res/train.pkl")
-    val_data = utils.load_from_pickle("./data/high_res/val.pkl")
-    test_data = utils.load_from_pickle("./data/high_res/test.pkl")
-
-    # process the data to keep only three channels
-    train_data = [(image[:, :, :3], label) for image, label in train_data]
-    val_data = [(image[:, :, :3], label) for image, label in val_data]
-    test_data = [(image[:, :, :3], label) for image, label in test_data]
-
-    # Create custom datasets with transformations
-    trainset = CustomDataset(train_data, transform=train_transform)
-    valset = CustomDataset(val_data, transform=transform)
-    testset = CustomDataset(test_data, transform=transform)
-
-    # Create DataLoaders
-    trainloader = DataLoader(trainset, batch_size=batch_size, shuffle=True)
-    valloader = DataLoader(valset, batch_size=batch_size, shuffle=False)
-    testloader = DataLoader(testset, batch_size=batch_size, shuffle=False)
-
-    # Set the device
-    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    print("Running on {}".format(device))
-
-    # Load a pretrained ResNet model
-    model = models.resnet50(weights="ResNet50_Weights.IMAGENET1K_V1").to(
-        device
-    )  # You can use resnet18, resnet50, etc.
-
-    # Modify the output layer directly to match binary classification
-    model.fc = nn.Sequential(
-        nn.Linear(model.fc.in_features, 128),
-        nn.ReLU(inplace=True),
-        nn.Linear(128, 1),  # 1 output unit for binary classification
-    ).to(
-        device
-    )  # Make sure the head is also on the correct device
-
-    # Print the model architecture
-    print(model)
-
-    if train:
-        # Define loss function and optimizer
-        criterion = nn.BCEWithLogitsLoss()  # Binary classification
-        optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
-        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.25)
-
-        # Train the model
-        optimal_accuracy = 0
-        patience = 10
-
-        for epoch in tqdm.tqdm(range(epochs)):
-            model.train()
-            running_loss = 0.0
-
-            for i, data in enumerate(trainloader, 0):
-                # Move inputs and labels to the correct device
-                inputs, labels = data
-                inputs, labels = (
-                    inputs.to(device),
-                    labels.to(device).float(),
-                )  # Ensure labels are float for BCEWithLogitsLoss
-
-                # Zero the parameter gradients
-                optimizer.zero_grad()
-
-                # Forward + backward + optimize
-                outputs = model(inputs)
-                loss = criterion(
-                    outputs, labels.unsqueeze(1)
-                )  # Match output shape (N, 1)
-                loss.backward()
-                optimizer.step()
-
-                running_loss += loss.item()
-                if i % 2000 == 1999:
-                    print(f"[{epoch + 1}, {i + 1}] loss: {running_loss / 2000:.3f}")
-                    running_loss = 0.0
-
-            scheduler.step()
-
-            # Validate the model
-            model.eval()
-            correct = 0
-            total = 0
-
-            with torch.no_grad():
-                for data in valloader:
-                    images, labels = data
-                    images, labels = images.to(device), labels.to(device).float()
-                    outputs = model(images)
-                    predicted = (
-                        torch.sigmoid(outputs) > 0.5
-                    ).float()  # Apply sigmoid and threshold at 0.5
-                    total += labels.size(0)
-                    correct += (predicted == labels.unsqueeze(1)).sum().item()
-
-            accuracy = 100 * correct / total
-            print(f"Validation accuracy: {accuracy:.2f}%")
-
-            # early stopping
-            if accuracy > optimal_accuracy:
-                optimal_accuracy = accuracy
-                optimal_model = model.state_dict()
-                patience = 10
-            else:
-                patience -= 1
-
-            if patience == 0:
-                print("Early stopping")
-                break
-
-        print("Finished Training")
-
-        # Save the model
-        torch.save(optimal_model, "./weights/ResNet50.pth")
-        print("Model saved")
-
-    # Load the model and move it to the correct device
-    model.load_state_dict(
-        torch.load("./weights/ResNet50.pth", map_location=device, weights_only=True)
-    )
-    model.to(device)
-
-    # Test the model
-    model.eval()
-    correct = 0
-    total = 0
-    y_pred = []
-    y_true = []
-
-    with torch.no_grad():
-        for data in tqdm.tqdm(testloader):
-            images, labels = data
-            images, labels = images.to(device), labels.to(device).float()
-            outputs = model(images)
-            predicted = (
-                torch.sigmoid(outputs) > 0.5
-            ).float()  # Apply sigmoid and threshold at 0.5
-
-            # store predictions for CM
-            y_pred.extend(predicted.cpu().numpy())
-            y_true.extend(labels.cpu().numpy())
-
-            total += labels.size(0)
-            correct += (predicted == labels.unsqueeze(1)).sum().item()
-
-    accuracy = 100 * correct / total
-    print(f"Test accuracy: {accuracy:.2f}%")
-
-    print("Finished Testing")
-
-    # Confusion matrix
-    cm = confusion_matrix(y_true, y_pred, normalize="true")
-    plt.figure(figsize=(8, 8))
-    plt.imshow(cm, interpolation="nearest", cmap=plt.cm.Blues)
-    plt.title("Confusion Matrix")
-    plt.colorbar()
-    tick_marks = np.arange(len(classes))
-    plt.xticks(tick_marks, classes, rotation=45)
-    plt.yticks(tick_marks, classes)
-    plt.xlabel("Predicted")
-    plt.ylabel("True")
-    plt.savefig("confusion_matrix.pdf", format="pdf")
-    plt.show()
-
-    # Show test images with predicted label and actual label
-    def imshow(img, title=None):
-        """This function plots a tensor"""
-        img = img / 2 + 0.5  # unnormalize
-        npimg = img.cpu().numpy()  # convert to numpy for display
-        plt.imshow(np.transpose(npimg, (1, 2, 0)))  # reshape to (H, W, C)
-        if title is not None:
-            plt.title(title)
-        plt.show()
-
-    def show_predictions(model, dataloader, device, classes):
-        """Function to show images with predicted and actual labels."""
-        model.eval()
-
-        with torch.no_grad():
-            for i, data in enumerate(dataloader):
-                images, labels = data
-                images, labels = images.to(device), labels.to(device).float()
-
-                # Forward pass to get predictions
-                outputs = model(images)
-                outputs = outputs.squeeze(dim=1)  # Ensure shape is [batch_size]
-                predicted = (torch.sigmoid(outputs) > 0.5).float()
-
-                # Plot each image with its predicted and actual labels
-                for j in range(images.size(0)):
-                    imshow(images[j].cpu())  # Unnormalize and plot image
-
-                    # Convert predictions and labels to text (formal/informal)
-                    pred_label = classes[int(predicted[j].item())]
-                    actual_label = classes[int(labels[j].item())]
-
-                    # Display the predicted and actual labels
-                    print(f"Predicted: {pred_label}, Actual: {actual_label}")
-
-                    # Optionally, stop after displaying N images
-                    if i * len(images) + j >= 20:  # Show 5 images, adjust as needed
-                        return
-
-    # Call the function to display images along with predictions
-    # show_predictions(model, testloader, device, classes)
-
-
-if __name__ == "__main__":
-    main()