Dog Breed Classification

from functools import partial
from pathlib import Path
import os
import warnings

from dotenv import load_dotenv
from PIL import Image, ImageFile
from torchvision import datasets
import cv2
import face_recognition
import matplotlib.cbook
warnings.filterwarnings("ignore", category=matplotlib.cbook.mplDeprecation)
import matplotlib.pyplot as pyplot
import matplotlib.image as mpimage
import matplotlib.patches as patches
import numpy
try:
    import pyttsx3
    SPEAKABLE = True
except ImportError:
    print("pyttsx3 not available")
    SPEAKABLE = False
import seaborn
import torch
import torchvision.models as models
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optimizer
import torchvision.transforms as transforms

from neurotic.tangles.data_paths import DataPathTwo
from neurotic.tangles.timer import Timer
from neurotic.constants.imagenet_map import imagenet

MODEL_PATH = DataPathTwo(folder_key="MODELS")

class DogPaths:
    """holds the paths to the dog images"""
    def __init__(self) -> None:
        self._main = None
        self._training = None
        self._testing = None
        self._validation = None
        self._breed_count = None
        load_dotenv()
        return

    @property
    def main(self) -> DataPathTwo:
        """The path to the main folder"""
        if self._main is None:
            self._main = DataPathTwo(folder_key="DOG_PATH")
        return self._main

    @property
    def training(self) -> DataPathTwo:
        """Path to the training images"""
        if self._training is None:
            self._training = DataPathTwo(folder_key="DOG_TRAIN")
        return self._training

    @property
    def validation(self) -> DataPathTwo:
        """Path to the validation images"""
        if self._validation is None:
            self._validation = DataPathTwo(folder_key="DOG_VALIDATE")
        return self._validation

    @property
    def testing(self) -> DataPathTwo:
        """Path to the testing images"""
        if self._testing is None:
            self._testing = DataPathTwo(folder_key="DOG_TEST")
        return self._testing

    @property
    def breed_count(self) -> int:
        """Counts the number of dog breeds

        This assumes that the training folder has all the breeds
        """
        if self._breed_count is None:
            self._breed_count = len(set(self.training.folder.iterdir()))
        return self._breed_count

    def check(self) -> None:
        """Checks that the folders are valid

        Raises: 
         AssertionError: folder doesn't exist
        """
        self.main.check_folder()
        self.training.check_folder()
        self.validation.check_folder()
        self.testing.check_folder()
        return

dog_paths = DogPaths()

human_path = DataPathTwo(folder_key="HUMAN_PATH")

print(dog_paths.main.folder)
print(dog_paths.training.folder)
print(dog_paths.testing.folder)
print(dog_paths.validation.folder)
dog_paths.check()
print(human_path.folder)
human_path.check_folder()

class Tee:
    """Save the input to a file and print it

    Args:
     log_name: name to give the log    
     directory_path: path to the directory for the file
    """
    def __init__(self, log_name: str, 
                 directory_name: str="../../../logs/dog-breed-classifier") -> None:
        self.directory_name = directory_name
        self.log_name = log_name
        self._path = None
        self._log = None
        return

    @property
    def path(self) -> Path:
        """path to the log-file"""
        if self._path is None:
            self._path = Path(self.directory_name).expanduser()
            assert self._path.is_dir()
            self._path = self._path.joinpath(self.log_name)
        return self._path

    @property
    def log(self):
        """File object to write log to"""
        if self._log is None:
            self._log = self.path.open("w", buffering=1)
        return self._log

    def __call__(self, line: str) -> None:
        """Writes to the file and stdout

        Args:
         line: text to emit
        """
        self.log.write("{}\n".format(line))
        print(line)
        return

 class F1Scorer:
     """Calculates the F1 and other scores

     Args:
      predictor: callable that gets passed and image and outputs boolean
      true_images: images that should be predicted as True
      false_images: images that shouldn't be matched by the predictor
      done_message: what to announce when done
     """
     def __init__(self, predictor: callable, true_images:list,
                  false_images: list,
                  done_message: str="Scoring Done") -> None:
         self.predictor = predictor
         self.true_images = true_images
         self.false_images = false_images
         self.done_message = done_message
         self._timer = None
         self._false_image_predictions = None
         self._true_image_predictions = None
         self._false_positives = None
         self._false_negatives = None
         self._true_positives = None
         self._true_negatives = None
         self._false_positive_rate = None
         self._precision = None
         self._recall = None
         self._f1 = None
         self._accuracy = None
         self._specificity = None
         return

     @property
     def timer(self) -> Timer:
         if self._timer is None:
             self._timer = Timer(message=self.done_message, emit=False)
         return self._timer

     @property
     def false_image_predictions(self) -> list:
         """Predictions made on the false-images"""
         if self._false_image_predictions is None:
             self._false_image_predictions = [self.predictor(str(image))
                                              for image in self.false_images]
         return self._false_image_predictions

     @property
     def true_image_predictions(self) -> list:
         """Predictions on the true-images"""
         if self._true_image_predictions is None:
             self._true_image_predictions = [self.predictor(str(image))
                                             for image in self.true_images]
         return self._true_image_predictions

     @property
     def true_positives(self) -> int:
         """count of correct positive predictions"""
         if self._true_positives is None:
             self._true_positives = sum(self.true_image_predictions)
         return self._true_positives

     @property
     def false_positives(self) -> int:
         """Count of incorrect positive predictions"""
         if self._false_positives is None:
             self._false_positives = sum(self.false_image_predictions)
         return self._false_positives

     @property
     def false_negatives(self) -> int:
         """Count of images that were incorrectly classified as negative"""
         if self._false_negatives is None:
             self._false_negatives = len(self.true_images) - self.true_positives
         return self._false_negatives

     @property
     def true_negatives(self) -> int:
         """Count of images that were correctly ignored"""
         if self._true_negatives is None:
             self._true_negatives = len(self.false_images) - self.false_positives
         return self._true_negatives

     @property
     def accuracy(self) -> float:
         """fraction of correct predictions"""
         if self._accuracy is None:
             self._accuracy = (
                 (self.true_positives + self.true_negatives)
                 /(len(self.true_images) + len(self.false_images)))
         return self._accuracy

     @property
     def precision(self) -> float:
         """True-Positive with penalty for false positives"""
         if self._precision is None:
             self._precision = self.true_positives/(
                 self.true_positives + self.false_positives)
         return self._precision

     @property
     def recall(self) -> float:
         """fraction of correct images correctly predicted"""
         if self._recall is None:
             self._recall = (
                 self.true_positives/len(self.true_images))
         return self._recall

     @property
     def false_positive_rate(self) -> float:
         """fraction of incorrect images predicted as positive"""
         if self._false_positive_rate is None:
             self._false_positive_rate = (
                 self.false_positives/len(self.false_images))
         return self._false_positive_rate

     @property
     def specificity(self) -> float:
         """metric for how much to believe a negative prediction

         Specificity is 1 - false positive rate so you only need one or the other
         """
         if self._specificity is None:
             self._specificity = self.true_negatives/(self.true_negatives
                                                      + self.false_positives)
         return self._specificity

     @property
     def f1(self) -> float:
         """Harmonic Mean of the precision and recall"""
         if self._f1 is None:
             TP = 2 * self.true_positives
             self._f1 = (TP)/(TP + self.false_negatives + self.false_positives)
         return self._f1

     def __call__(self) -> None:
         """Emits the F1 and other scores as an org-table
         """
         self.timer.start()
         print("|Metric|Value|")
         print("|-+-|")
         print("|Accuracy|{:.2f}|".format(self.accuracy))
         print("|Precision|{:.2f}|".format(self.precision))
         print("|Recall|{:.2f}|".format(self.recall))
         print("|Specificity|{:.2f}".format(self.specificity))
         # print("|False Positive Rate|{:.2f}|".format(self.false_positive_rate))
         print("|F1|{:.2f}|".format(self.f1))
         self.timer.end()
         print("|Elapsed|{}|".format(self.timer.ended - self.timer.started))
         return

def get_name(path: Path) -> str:
    """Extracts the name of the person from the file name

    Args:
     path: path to the file

    Returns:
     the name extracted from the file name
    """
    return " ".join(path.name.split("_")[:-1]).title()

def display_image(image: Path, title: str, is_file: bool=True) -> tuple:
    """Plot the image

    Args:
     image: path to the image file or image
     title: title for the image
     is_file: first argument is a file name, not an array

    Returns:
     figure, axe
    """
    figure, axe = pyplot.subplots()
    figure.suptitle(title, weight="bold")
    axe.tick_params(dict(axis="both",
                         which="both",
                         bottom=False,
                         top=False))
    axe.get_xaxis().set_ticks([])
    axe.get_yaxis().set_ticks([])
    if is_file:
        image = Image.open(image)
    image = axe.imshow(image)
    return figure, axe

def first_prediction(source: list, start:int=0) -> int:
    """Gets the index of the first True prediction

    Args:
     source: list of True/False predictions
     start: index to start the search from

    Returns:
     index of first True prediction found
    """
    for index, prediction in enumerate(source[start:]):
        if prediction:
            print("{}: {}".format(start + index, prediction))
            break
    return start + index

haar_path = DataPathTwo("haarcascade_frontalface_alt.xml", folder_key="HAAR_CASCADES")
print(haar_path.from_folder)
assert haar_path.from_folder.is_file()

face_cascade = cv2.CascadeClassifier(str(haar_path.from_folder))

gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

faces = face_cascade.detectMultiScale(gray)

print('Number of faces detected:', len(faces))

for (x,y,w,h) in faces:
    # add bounding box to color image
    cv2.rectangle(image, (x,y), (x+w,y+h), (255,0,0), 2)

cv_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

figure, axe = display_image(cv_rgb, "OpenCV Face-Detection Bounding Box", False)

# returns "True" if face is detected in image stored at img_path
def face_detector(image_path: str) -> bool:
    """Detects human faces in an image

    Args:
     image_path: path to the image to check

    Returns:
     True if there was at least one face in the image
    """
    image = cv2.imread(image_path)
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(gray)
    return len(faces) > 0

open_cv_scorer = human_scorer(face_detector)
open_cv_scorer()

dogman_index = first_prediction(open_cv_scorer.false_image_predictions)

source = dog_files_short[dogman_index]
name = get_name(source)
figure, axe = display_image(source,
                            "Dog-Human OpenCV Prediction ({})".format(name))

image = face_recognition.load_image_file(sample_face)
locations = face_recognition.face_locations(image)
image = mpimage.imread(sample_face)

figure, axe = display_image(image, "dlib Face Recognition Bounding-Box", False)
top, right, bottom, left = locations[0]
width = right - left
height = top - bottom
rectangle = patches.Rectangle((top, right), width, height, fill=False)
patch = axe.add_patch(rectangle)

def face_recognition_check(image_path: str) -> bool:
    """This decides if an image has a face in it

    Args:
     image_path: path to an image
    Returns:
     True if there's at least one face in the image
    """
    image = face_recognition.load_image_file(str(image_path))
    locations = face_recognition.face_locations(image)
    return len(locations) > 0

dlib_dog_humans = human_scorer(face_recognition_check)
dlib_dog_humans()

dlib_dog_human_index = first_prediction(dlib_dog_humans.false_image_predictions)

source = dog_files_short[dlib_dog_human_index]
name = get_name(source)
figure, axe = display_image(source,
                            "Dog-Human DLib Prediction ({})".format(name))

IMAGE_SIZE = 224
IMAGE_HALF_SIZE = IMAGE_SIZE//2

vgg_transform = transforms.Compose([transforms.Resize(255),
                                    transforms.CenterCrop(IMAGE_SIZE),
                                    transforms.ToTensor(),
                                    transforms.Normalize(MEANS,
                                                         DEVIATIONS)])

def VGG16_predict(img_path: str) -> int:
    '''
    Uses a pre-trained VGG-16 model to obtain the index corresponding to 
    predicted ImageNet class for image at specified path

    Args:
        img_path: path to an image

    Returns:
        Index corresponding to VGG-16 model's prediction
    '''
    image = Image.open(str(img_path))
    image = vgg_transform(image).unsqueeze(0).to(device)
    output = VGG16(image)
    probabilities = torch.exp(output)
    top_probability, top_class = probabilities.topk(1, dim=1)
    return top_class.item()

path = numpy.random.choice(dog_files_short)
print(path)
classification = VGG16_predict(path)
print(imagenet[classification])

name = get_name(path)
figure, axe = display_image(path, name)

 def inception_predicts(image_path: str) -> int:
     """Predicts the category of the image

     Args:
      image_path: path to the image file

     Returns:
      classification: the resnet ID for the image
     """
     image = Image.open(str(image_path))
     image = vgg_transform(image).unsqueeze(0).to(device)
     output = inception(image)
     probabilities = torch.exp(output)
     top_probability, top_class = probabilities.topk(1, dim=1)
     return top_class.item()

 for path in dog_files_short:
     try:
         prediction = inception_predicts(path)
     except RuntimeError as error:
         print(error)
         print(path)
         break

 name = get_name(path)
 figure, axe = display_image(path, "Error-Producing Image ({})".format(name))

inception_predictor = partial(dog_detector, predictor=inception_predicts_two)
inception_scorer = dog_scorer(inception_predictor)
inception_scorer()

inception_human_dog = first_prediction(inception_scorer.false_image_predictions)

figure, axe = pyplot.subplots()
source = human_files_short[inception_human_dog]
name = " ".join(
    os.path.splitext(
        os.path.basename(source))[0].split("_")[:-1]).title()
figure.suptitle("Human-Dog Inception Prediction ({})".format(
    name), weight="bold")
image = Image.open(source)
image = axe.imshow(image)

IMAGE_SIZE = 224
IMAGE_HALF_SIZE = IMAGE_SIZE//2

train_transform = transforms.Compose([
    transforms.RandomRotation(30),
    transforms.RandomResizedCrop(IMAGE_SIZE),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(MEANS,
                         DEVIATIONS)])

test_transform = transforms.Compose([transforms.Resize(255),
                                      transforms.CenterCrop(IMAGE_SIZE),
                                      transforms.ToTensor(),
                                      transforms.Normalize(MEANS,
                                                           DEVIATIONS)])

training = datasets.ImageFolder(root=str(dog_paths.training.folder),
                                transform=train_transform)
validation = datasets.ImageFolder(root=str(dog_paths.validation.folder),
                                  transform=test_transform)
testing = datasets.ImageFolder(root=str(dog_paths.testing.folder),
                               transform=test_transform)

BATCH_SIZE = 35
WORKERS = 0

train_batches = torch.utils.data.DataLoader(training, batch_size=BATCH_SIZE,
                                            shuffle=True, num_workers=WORKERS)
validation_batches = torch.utils.data.DataLoader(
    validation, batch_size=BATCH_SIZE, shuffle=True, num_workers=WORKERS)
test_batches = torch.utils.data.DataLoader(
    testing, batch_size=BATCH_SIZE, shuffle=True, num_workers=WORKERS)

loaders_scratch = dict(train=train_batches,
                       validate=validation_batches,
                       test=test_batches)

LAYER_ONE_OUT = 16
LAYER_TWO_OUT = LAYER_ONE_OUT * 2
LAYER_THREE_OUT = LAYER_TWO_OUT * 2

KERNEL = 3
PADDING = 1
FULLY_CONNECTED_OUT = 500

conv_1 = nn.Conv2d(3, LAYER_ONE_OUT, KERNEL, padding=PADDING)
conv_2 = nn.Conv2d(LAYER_ONE_OUT, LAYER_TWO_OUT, KERNEL, padding=PADDING)
conv_3 = nn.Conv2d(LAYER_TWO_OUT, LAYER_THREE_OUT, KERNEL, padding=PADDING)

pool = nn.MaxPool2d(2, 2)
dropout = nn.Dropout(0.25)

fully_connected_1 = nn.Linear((IMAGE_HALF_SIZE//4)**2 * LAYER_THREE_OUT, FULLY_CONNECTED_OUT)
fully_connected_2 = nn.Linear(FULLY_CONNECTED_OUT, dog_paths.breed_count)

dataiter = iter(loaders_scratch['train'])
images, labels = dataiter.next()

x = pool(F.relu(conv_1(images)))
print(x.shape)
assert x.shape == torch.Size([BATCH_SIZE, 16, IMAGE_HALF_SIZE, IMAGE_HALF_SIZE])

x = pool(F.relu(conv_2(x)))
print(x.shape)
assert x.shape == torch.Size([BATCH_SIZE, LAYER_TWO_OUT, IMAGE_HALF_SIZE//2, IMAGE_HALF_SIZE//2])

x = pool(F.relu(conv_3(x)))
print(x.shape)
assert x.shape == torch.Size([BATCH_SIZE, LAYER_THREE_OUT, IMAGE_HALF_SIZE//4, IMAGE_HALF_SIZE//4])

x = x.view(-1, ((IMAGE_HALF_SIZE//4)**2) * LAYER_THREE_OUT)
print(x.shape)
x = fully_connected_1(x)
print(x.shape)
x = fully_connected_2(x)
print(x.shape)

class NaiveNet(nn.Module):
    """Naive Neural Network to classify dog breeds"""
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(3, LAYER_ONE_OUT,
                               KERNEL, padding=PADDING)
        self.conv2 = nn.Conv2d(LAYER_ONE_OUT, LAYER_TWO_OUT,
                               KERNEL, padding=PADDING)
        self.conv3 = nn.Conv2d(LAYER_TWO_OUT, LAYER_THREE_OUT,
                               KERNEL, padding=PADDING)
        # max pooling layer
        self.pool = nn.MaxPool2d(2, 2)
        # linear layer
        self.fc1 = nn.Linear((IMAGE_HALF_SIZE//4)**2 * LAYER_THREE_OUT, FULLY_CONNECTED_OUT)
        self.fc2 = nn.Linear(FULLY_CONNECTED_OUT, BREEDS)
        # dropout layer (p=0.25)
        self.dropout = nn.Dropout(0.25)
        return


    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """The forward pass method

        Args:
         x: a n x 224 x 224 x 3 tensor

        Returns:
         tensor of probabilities
        """
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))

        x = x.view(-1, (IMAGE_HALF_SIZE//4)**2 * LAYER_THREE_OUT)
        x = self.dropout(x)

        x = self.dropout(F.relu(self.fc1(x)))
        x = self.fc2(x)        
        return x

naive_model = NaiveNet()
naive_model.to(device)
MODELS.append(naive_model)

criterion_scratch = nn.CrossEntropyLoss()

optimizer_scratch = optimizer.SGD(naive_model.parameters(),
                                  lr=0.001,
                                  momentum=0.9)

OSError: image file is truncated (150 bytes not processed)

timer.start()
broken = None
for image in dog_files:
    try:
        opened = Image.open(image)
        opened.convert("RGB")
    except OSError as error:
        print("{}: {}".format(error, image))
        broken = image
timer.end()

figure, axe = pyplot.subplots()
name = " ".join(broken.name.split("_")[:-1]).title()
figure.suptitle("Truncated Image ({})".format(name), weight="bold")
image = Image.open(broken)
axe_image = axe.imshow(image)

ImageFile.LOAD_TRUNCATED_IMAGES = True

NAIVE_PATH = MODEL_PATH.folder.joinpath("model_scratch.pt")
scratch_log = Tee(log_name="scratch_train.log")

def test(test_batches: torch.utils.data.DataLoader,
         model: nn.Module,
         criterion: nn.CrossEntropyLoss) -> None:
    """Test the model

    Args:
     test_batches: batch loader of test images
     model: the network to test
     criterion: calculator for the loss
    """
    test_loss = 0.
    correct = 0.
    total = 0.

    model.eval()
    for data, target in test_batches:
        data, target = data.to(device), target.to(device)
        output = model(data)
        loss = criterion(output, target)
        test_loss += loss.item() * data.size(0)
        # convert output probabilities to predicted class
        predictions = output.data.max(1, keepdim=True)[1]
        # compare predictions to true label
        correct += numpy.sum(
            numpy.squeeze(
                predictions.eq(
                    target.data.view_as(predictions))).cpu().numpy())
        total += data.size(0)
    test_loss /= len(test_batches.dataset)
    print('Test Loss: {:.6f}\n'.format(test_loss))
    print('\nTest Accuracy: %2d%% (%2d/%2d)' % (
        100. * correct / total, correct, total))
    return

model_scratch.load_state_dict(torch.load('model_scratch.pt'))

test(loaders_scratch["test"], model_scratch, criterion_scratch)

class Transformer:
    """builds the data-sets

    Args:
     means: list of means for each channel
     deviations: list of standard deviations for each channel
     image_size: size to crop the image to
    """
    def __init__(self,
                 means: list=[0.485, 0.456, 0.406],
                 deviations: list=[0.229, 0.224, 0.225],
                 image_size: int=299) -> None:
        self.means = means
        self.deviations = deviations
        self.image_size = image_size
        self._training = None
        self._testing = None
        return

    @property
    def training(self) -> transforms.Compose:
        """The image transformers for the training"""
        if self._training is None:
            self._training = transforms.Compose([
                transforms.RandomRotation(30),
                transforms.RandomResizedCrop(self.image_size),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize(self.means,
                                     self.deviations)])
        return self._training

    @property
    def testing(self) -> transforms.Compose:
        """Image transforms for the testing"""
        if self._testing is None:
            self._testing = transforms.Compose(
                [transforms.Resize(350),
                 transforms.CenterCrop(self.image_size),
                 transforms.ToTensor(),
                 transforms.Normalize(self.means,
                                      self.deviations)])
        return self._testing

class DataSets:
    """Builds the data-sets

    Args:
     paths: object with the paths to the data-sets
    """
    def __init__(self, paths: DogPaths=None, transformer: Transformer=None) -> None:
        self._paths = paths
        self._transformer = transformer
        self._training = None
        self._validation = None
        self._testing = None
        return

    @property
    def paths(self) -> DogPaths:
        """Object with the paths to the image files"""
        if self._paths is None:
            self._paths = DogPaths()
        return self._paths

    @property
    def transformer(self) -> Transformer:
        """Object with the image transforms"""
        if self._transformer is None:
            self._transformer = Transformer()
        return self._transformer

    @property
    def training(self) -> datasets.ImageFolder:
        """The training data set"""
        if self._training is None:
            self._training = datasets.ImageFolder(
                root=self.paths.training.folder,
                transform=self.transformer.training)
        return self._training

    @property
    def validation(self) -> datasets.ImageFolder:
        """The validation dataset"""
        if self._validation is None:
            self._validation = datasets.ImageFolder(
                root=self.paths.validation.folder,
                transform=self.transformer.testing)
        return self._validation

    @property
    def testing(self) -> datasets.ImageFolder:
        """The test set"""
        if self._testing is None:
            self._testing = datasets.ImageFolder(
                root=self.paths.testing.folder,
                transform=self.transformer.testing)
        return self._testing

class Batches:
    """The data batch loaders

    Args:
     datasets: a data-set builder
     batch_size: the size of each batch loaded
     workers: the number of processes to use
    """
    def __init__(self, datasets: DataSets,
                 batch_size: int=20,
                 workers: int=0) -> None:
        self.datasets = datasets
        self.batch_size = batch_size
        self.workers = workers
        self._training = None
        self._validation = None
        self._testing = None
        return

    @property
    def training(self) -> torch.utils.data.DataLoader:
        """The training batches"""
        if self._training is None:
            self._training = torch.utils.data.DataLoader(
                self.datasets.training,
                batch_size=self.batch_size,
                shuffle=True, num_workers=self.workers)
        return self._training

    @property
    def validation(self) -> torch.utils.data.DataLoader:
        """The validation batches"""
        if self._validation is None:
            self._validation = torch.utils.data.DataLoader(
                self.datasets.validation,
                batch_size=self.batch_size,
                shuffle=True, num_workers=self.workers)
        return self._validation

    @property
    def testing(self) -> torch.utils.data.DataLoader:
        """The testing batches"""
        if self._testing is None:
            self._testing = torch.utils.data.DataLoader(
                self.datasets.testing,
                batch_size=self.batch_size,
                shuffle=True, num_workers=self.workers)
        return self._testing

class Inception:
    """Sets up the model, criterion, and optimizer for the transfer learning

    Args:
     classes: number of outputs for the final layer
     device: processor to use
     model_path: path to a saved model
     learning_rate: learning rate for the optimizer
     momentum: momentum for the optimizer
    """
    def __init__(self, classes: int,
                 device: torch.device=None,
                 model_path: str=None,
                 learning_rate: float=0.001, momentum: float=0.9) -> None:
        self.classes = classes
        self.model_path = model_path
        self.learning_rate = learning_rate
        self.momentum = momentum
        self._device = device
        self._model = None
        self._classifier_inputs = None
        self._criterion = None
        self._optimizer = None
        return

    @property
    def device(self) -> torch.device:
        """Processor to use (cpu or cuda)"""
        if self._device is None:
            self._device = torch.device(
                "cuda" if torch.cuda.is_available() else "cpu")
        return self._device

    @property
    def model(self) -> models.inception_v3:
        """The inception model"""
        if self._model is None:
            self._model = models.inception_v3(pretrained=True)
            for parameter in self._model.parameters():
                parameter.requires_grad = False
            classifier_inputs = self._model.fc.in_features
            self._model.fc = nn.Linear(in_features=classifier_inputs,
                                       out_features=self.classes,
                                       bias=True)
            self._model.to(self.device)
            if self.model_path:
                self._model.load_state_dict(torch.load(self.model_path))
        return self._model

    @property
    def criterion(self) -> nn.CrossEntropyLoss:
        """The loss callable"""
        if self._criterion is None:
            self._criterion = nn.CrossEntropyLoss()
        return self._criterion

    @property
    def optimizer(self) -> optimizer.SGD:
        """The Gradient Descent object"""
        if self._optimizer is None:
            self._optimizer = optimizer.SGD(
                self.model.parameters(),
                lr=self.learning_rate,
                momentum=self.momentum)
        return self._optimizer

# python
from pathlib import Path
from functools import partial

import argparse

# pypi
from dotenv import load_dotenv
from PIL import ImageFile
from torchvision import datasets
import numpy
import torch
import torch.nn as nn
import torch.optim as optimizer
import torchvision.models as models
import torchvision.transforms as transforms

# this project
from neurotic.tangles.data_paths import DataPathTwo
from neurotic.tangles.timer import Timer

# the output won't show up if you don't flush it when redirecting it to a file
print = partial(print, flush=True)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Test or Train the Inception V3 Dog Classifier")
    parser.add_argument("--test-only", action="store_true",
                        help="Only run the test")
    parser.add_argument("--epochs", default=10, type=int,
                        help="Training epochs (default: %(default)s)")
    parser.add_argument(
        "--epoch-offset", default=1, type=int,
        help="Offset for the output of epochs (default: %(default)s)")
    parser.add_argument("--restart", action="store_true",
                        help="Wipe out old model.")

    arguments = parser.parse_args()

    data_sets = DataSets(training_path=dog_training_path.folder,
                         validation_path=dog_validation_path.folder,
                         testing_path=dog_testing_path.folder)
    batches = Batches(datasets=data_sets)
    inception = Inception(classes=len(data_sets.training.classes)
    trainer = Trainer(epochs=arguments.epochs,
                      epoch_start=arguments.epoch_offset,
                      training_batches=batches.training,
                      validation_batches=batches.validation,
                      testing_batches=batches.testing,
                      model=inception.model,
                      device=inception.device,
                      optimizer=inception.optimizer,
                      criterion=inception.criterion,
                      model_path=transfer_path.from_folder,
                      load_model=True,
                      beep=False)
    if arguments.test_only:
        trainer.test()
    else:
        trainer()

transfer_path = MODEL_PATH.folder.joinpath("model_transfer.pt")
transfer_log = Tee(log_name="transfer_train.log")

EPOCHS = 100
inception = Inception()
train(EPOCHS,
      loaders=loaders_transfer,
      model=inception.model,
      optimizer=inception.optimizer,
      criterion=inception.criterion,
      use_cuda=use_cuda,
      save_path=transfer_model_path,
      print_function=transfer_log,
      is_inception=True)

model_transfer.load_state_dict(torch.load(transfer_model_path))
transfer_test_log = Tee("transfer_test.log")
test(loaders_transfer, model_transfer, criterion_transfer, use_cuda, print_function=transfer_test_log)