# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/01_DefineModels.ipynb (unless otherwise specified).

__all__ = ['data_transforms', 'Flatten', 'deeptracknet2']

# Cell
import torch
from torchvision import transforms

data_transforms = transforms.Compose([
    transforms.Grayscale(1),
    transforms.Resize(51),
    transforms.ToTensor()
])

# Cell
from torch import nn

class Flatten(nn.Module):
    def forward(self, x):
        return x.view(x.size()[0], -1)

deeptracknet2 = nn.Sequential(
        nn.Conv2d(1,16,3),
        nn.ReLU(),
        nn.MaxPool2d(2),
        nn.Conv2d(16,32,3),
        nn.ReLU(),
        nn.MaxPool2d(2),
        nn.Conv2d(32,64,3),
        nn.ReLU(),
        nn.MaxPool2d(2),
        Flatten(),
        nn.Linear(1024,32),
        nn.ReLU(),
        nn.Linear(32,32),
        nn.ReLU(),
        nn.Linear(32,2)
    )