gtsrb_aug.py

""" Models obtained from https://github.com/poojahira/gtsrb-pytorch """
from __future__ import print_function
import zipfile
import os

import torchvision.transforms as transforms

# data augmentation for training and test time
# Resize all images to 32 * 32 and normalize them to mean = 0 and standard-deviation = 1 based on statistics collected from the training set

data_transforms = transforms.Compose(
    [
        transforms.Resize((32, 32)),
        transforms.ToTensor(),
    ]
)

# Resize, normalize and jitter image brightness
data_jitter_brightness = transforms.Compose(
    [
        transforms.Resize((32, 32)),
        transforms.ColorJitter(brightness=5),
        transforms.ToTensor(),
    ]
)

# Resize, normalize and jitter image saturation
data_jitter_saturation = transforms.Compose(
    [
        transforms.Resize((32, 32)),
        transforms.ColorJitter(saturation=5),
        transforms.ToTensor(),
    ]
)

# Resize, normalize and jitter image contrast
data_jitter_contrast = transforms.Compose(
    [
        transforms.Resize((32, 32)),
        transforms.ColorJitter(contrast=5),
        transforms.ToTensor(),
    ]
)

# Resize, normalize and jitter image hues
data_jitter_hue = transforms.Compose(
    [
        transforms.Resize((32, 32)),
        transforms.ColorJitter(hue=0.4),
        transforms.ToTensor(),
    ]
)

# Resize, normalize and rotate image
data_rotate = transforms.Compose(
    [
        transforms.Resize((32, 32)),
        transforms.RandomRotation(15),
        transforms.ToTensor(),
    ]
)

# Resize, normalize and flip image horizontally and vertically
data_hvflip = transforms.Compose(
    [
        transforms.Resize((32, 32)),
        transforms.RandomHorizontalFlip(1),
        transforms.RandomVerticalFlip(1),
        transforms.ToTensor(),
    ]
)

# Resize, normalize and flip image horizontally
data_hflip = transforms.Compose(
    [
        transforms.Resize((32, 32)),
        transforms.RandomHorizontalFlip(1),
        transforms.ToTensor(),
    ]
)

# Resize, normalize and flip image vertically
data_vflip = transforms.Compose(
    [
        transforms.Resize((32, 32)),
        transforms.RandomVerticalFlip(1),
        transforms.ToTensor(),
    ]
)

# Resize, normalize and shear image
data_shear = transforms.Compose(
    [
        transforms.Resize((32, 32)),
        transforms.RandomAffine(degrees=15, shear=2),
        transforms.ToTensor(),
    ]
)

# Resize, normalize and translate image
data_translate = transforms.Compose(
    [
        transforms.Resize((32, 32)),
        transforms.RandomAffine(degrees=15, translate=(0.1, 0.1)),
        transforms.ToTensor(),
    ]
)

# Resize, normalize and crop image
data_center = transforms.Compose(
    [
        transforms.Resize((36, 36)),
        transforms.CenterCrop(32),
        transforms.ToTensor(),
    ]
)

# Resize, normalize and convert image to grayscale
data_grayscale = transforms.Compose(
    [
        transforms.Resize((32, 32)),
        transforms.Grayscale(num_output_channels=3),
        transforms.ToTensor(),
    ]
)


def initialize_data(folder):
    train_zip = folder + "/train_images.zip"
    test_zip = folder + "/test_images.zip"
    if not os.path.exists(train_zip) or not os.path.exists(test_zip):
        raise (
            RuntimeError(
                "Could not find "
                + train_zip
                + " and "
                + test_zip
                + ", please download them from https://www.kaggle.com/c/nyu-cv-fall-2017/data "
            )
        )

    # extract train_data.zip to train_data
    train_folder = folder + "/train_images"
    if not os.path.isdir(train_folder):
        print(train_folder + " not found, extracting " + train_zip)
        zip_ref = zipfile.ZipFile(train_zip, "r")
        zip_ref.extractall(folder)
        zip_ref.close()

    # extract test_data.zip to test_data
    test_folder = folder + "/test_images"
    if not os.path.isdir(test_folder):
        print(test_folder + " not found, extracting " + test_zip)
        zip_ref = zipfile.ZipFile(test_zip, "r")
        zip_ref.extractall(folder)
        zip_ref.close()

    # make validation_data by using images 00000*, 00001* and 00002* in each class
    val_folder = folder + "/val_images"
    if not os.path.isdir(val_folder):
        print(val_folder + " not found, making a validation set")
        os.mkdir(val_folder)
        for dirs in os.listdir(train_folder):
            if dirs.startswith("000"):
                os.mkdir(val_folder + "/" + dirs)
                for f in os.listdir(train_folder + "/" + dirs):
                    if (
                        f.startswith("00000")
                        or f.startswith("00001")
                        or f.startswith("00002")
                    ):
                        # move file to validation folder
                        os.rename(
                            train_folder + "/" + dirs + "/" + f,
                            val_folder + "/" + dirs + "/" + f,
                        )