import torch import torch.nn as nn import torch.nn.functional as F from torchvision import transforms from torch.utils.data import DataLoader from datasets import load_dataset from PIL import Image import matplotlib.pyplot as plt import numpy as np import math from tqdm import tqdm # creates image and mask transforms. Can discuss hyperparameters later, but we have 256 x 256 images, and normalize to [-1, 1] IMG_SIZE = 256 image_transform = transforms.Compose([ transforms.Resize((IMG_SIZE, IMG_SIZE)), transforms.ToTensor(), transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), ]) mask_transform = transforms.Compose([ transforms.Resize((IMG_SIZE, IMG_SIZE), interpolation=transforms.InterpolationMode.NEAREST), transforms.PILToTensor(), ]) # Helpful methods for visualizing the images and masks def visualize_mask(mask: Image, img: Image = None): """ Visualizes the segmentation mask. If an image is provided, it overlays the mask on the image. :param mask: The segmentation mask to visualize. Expects a pillow image :param img: The image to overlay the mask on. Expects a pillow image :return: """ mask_np = np.array(mask) class_colors = { 1: [255, 0, 0], # red for prediction 2: [0, 255, 0], # green for background 3: [0, 0, 255], # blue for ambiguous } h, w = mask_np.shape color_mask = np.zeros((h, w, 3), dtype=np.uint8) for class_id, color in class_colors.items(): color_mask[mask_np == class_id] = color if img is not None: img = img.convert('RGBA') overlay = Image.fromarray(color_mask).convert('RGBA') blended = Image.blend(img, overlay, alpha=0.5) plt.figure(figsize=(6, 6)) plt.imshow(blended) plt.title('Segmentation Overlay') plt.axis('off') plt.show() else: plt.figure(figsize=(6, 6)) plt.imshow(color_mask) plt.title('Colorized Segmentation Mask') plt.axis('off') plt.show() def visualize_mask_tensor(img_tensor: torch.Tensor, mask_tensor: torch.Tensor, alpha: float = 0.5, title: str = "Image Mask"): """ Overlays a mask tensor onto an image tensor and displays the result. :param img_tensor: Image tensor of shape (C, H, W) :param mask_tensor: Mask tensor of shape (H, W) :param mode: 'class' or 'category' :param alpha: Transparency for overlay """ img_np = img_tensor.permute(1, 2, 0).numpy() mask_np = mask_tensor.numpy() h, w = mask_np.shape color_mask = np.zeros((h, w, 3), dtype=np.uint8) color_mask[mask_np == 1] = [255, 0, 0] # red for mask color_mask[mask_np == 2] = [0, 0, 255] # blue for bg plt.figure(figsize=(6, 6)) plt.imshow(img_np) plt.imshow(color_mask, alpha=alpha) plt.axis('off') plt.title(title) plt.show() class_labels = ['cat', 'dog'] category_labels = ['Abyssinian', 'american bulldog', 'american pit bull terrier', 'basset hound', 'beagle', 'Bengal', 'Birman', 'Bombay', 'boxer', 'British Shorthair', 'chihuahua', 'Egyptian Mau', 'english cocker spaniel', 'english setter', 'german shorthaired', 'great pyrenees', 'havanese', 'japanese chin', 'keeshond', 'leonberger', 'Maine Coon', 'miniature pinscher', 'newfoundland', 'Persian', 'pomeranian', 'pug', 'Ragdoll', 'Russian Blue', 'saint bernard', 'samoyed', 'scottish terrier', 'shiba inu', 'Siamese', 'Sphynx', 'staffordshire bull terrier', 'wheaten terrier', 'yorkshire terrier'] # the index of these lists relate to the label that is associate with an image def transform_class_example(example: dict, include_ambiguous: bool = False) -> dict: """ Transforms the image and mask in from a given example :param example: an example from a hf dataset; a dictionary with keys from the dataset :param include_ambiguous: whether to include the ambiguous class in the mask. If true, the ambiguous class is included in the mask. If false, the ambiguous class is removed from the mask (background). :return: a dictionary with the transformed image and mask """ class_label = example['class'] mask = mask_transform(example["msk"]).squeeze(0) mask[mask == 0] = 2 mask[mask == 1] = 0 mask[mask == 2] = 1 if include_ambiguous: mask[mask == 3] = 1 else: mask[mask == 3] = 0 return { "image": image_transform(example["img"]), "mask": mask, "classification": class_label } def transform_category_example(example: dict, include_ambiguous: bool = False) -> dict: """ Transforms the image and mask in from a given example :param example: an example from a hf dataset; a dictionary with keys from the dataset :param include_ambiguous: whether to include the ambiguous class in the mask. If true, the ambiguous class is included in the mask. If false, the ambiguous class is removed from the mask (background). :return: a dictionary with the transformed image and mask """ category_label = example['category'] mask = mask_transform(example["msk"]).squeeze(0) # switch 0 and 1. Now 0 is background and 1 is the mask mask[mask == 0] = 2 mask[mask == 1] = 0 mask[mask == 2] = 1 if include_ambiguous: mask[mask == 3] = 1 else: mask[mask == 3] = 0 # here we are switching around the values from the original dataset. Instead of having 0 be foreground, it becomes the vaue # of the category number, the background becomes 37 and ambiguous becomes 38 return { "image": image_transform(example["img"]), "mask": mask, "classification": category_label } def generate_dataset(split: str, classification: str = 'class'): """ Generates a dataset for the given split and classification type. :param split: which split to generate. Can be train, test or valid :param mask: decides whether to use the mask, or the class label :param classification: Can either be class or category. Class is either cat or dog, while category is the breed. :return: a dataset of the given split. With transformed images and masks. """ if split not in ['test', 'train', 'valid']: raise ValueError('Split must be either test train or valid') if classification not in ['class', 'category']: raise ValueError('Classification must be either class or category. Class is either cat or dog, while category is the breed.') dataset = load_dataset('cvdl/oxford-pets', split=split) # remove the bbox and non-used classification columns if classification == 'class': dataset = dataset.map(transform_class_example, remove_columns=['bbox', 'img', 'msk', 'category', 'class']) elif classification == 'category': dataset = dataset.map(transform_category_example, remove_columns=['bbox', 'img', 'msk', 'class', 'category']) dataset.set_format(type='torch', columns=['image', 'mask', 'classification']) # transform the images and masks return dataset train_set = generate_dataset('train', classification='class') test_set = generate_dataset('test', classification='class') mask = True batch_size = 32 train_loader = DataLoader(train_set, batch_size=batch_size) test_loader = DataLoader(test_set, batch_size=batch_size) class LeNet(nn.Module): def __init__(self, num_classes=2): super(LeNet, self).__init__() self.conv1 = nn.Conv2d(3, 6, 5, 1) self.conv2 = nn.Conv2d(6, 12, 5, 1) self.conv3 = nn.Conv2d(12, 24, 5, 1) self.lin1 = nn.Linear(24*28*28, 1024) self.lin2 = nn.Linear(1024, 512) self.lin3 = nn.Linear(512, 128) self.lin4 = nn.Linear(128, num_classes) def forward(self, x): x = nn.functional.max_pool2d(nn.functional.relu(self.conv1(x)), (2,2)) x = nn.functional.max_pool2d(nn.functional.relu(self.conv2(x)), (2,2)) x = nn.functional.max_pool2d(nn.functional.relu(self.conv3(x)), (2,2)) x = torch.flatten(x, 1) x = nn.functional.relu(self.lin1(x)) x = nn.functional.relu(self.lin2(x)) x = nn.functional.relu(self.lin3(x)) out = self.lin4(x) return out def train_model(model, train_loader, optimizer, criterion, epoch): model.train() train_loss = 0.0 for idx, batch in tqdm(enumerate(train_loader)): optimizer.zero_grad() # print(batch['image'].shape, batch['classification'].shape) output = model(batch['image']) loss = criterion(output, batch['classification']) loss.backward() optimizer.step() train_loss += loss.item() train_loss /= len(train_loader) print('[Training set] Epoch: {:d}, Average loss: {:.4f}'.format(epoch+1, train_loss)) return train_loss def test_model(model, test_loader, epoch): model.eval() correct = 0 with torch.no_grad(): for idx, value in tqdm(enumerate(test_loader)): output = model(value['image']) batch_pred = output.max(1, keepdim=True)[1] correct += batch_pred.eq(value['classification'].view_as(batch_pred)).sum().item() # print(correct, len(test_loader.dataset)) test_acc = correct / len(test_loader.dataset) print('[Test set] Epoch: {:d}, Accuracy: {:.2f}%\n'.format( epoch+1, 100. * test_acc)) return test_acc model = LeNet(num_classes=2) criterion = nn.CrossEntropyLoss() optimizer = torch.optim.Adam(model.parameters(), lr=0.0001) best_acc = 0 for epoch in range(0, 100): print("now training") train_model(model, train_loader, optimizer, criterion, epoch) print("now testing") acc = test_model(model, test_loader, epoch) print(f"epoch {epoch+1}, best_acc {max(best_acc, acc)}") best_acc = max(best_acc, acc) torch.save(model.state_dict(), f'animals_nn_epoch_{epoch+1}.pth')