working poc 🌭

mACHINE-LEARNINGS/classes/MemeDataset.py

@@ -1,24 +1,37 @@
 import os
-import pandas as pd
-from torchvision.io import read_image
+import glob
+from torchvision.transforms import Resize
+from torchvision.io import read_image, ImageReadMode
 from torch.utils.data import Dataset
-from sklearn import preprocessing
+import torch
+import imghdr
+
+supported_ext = ['jpg']
 
 class MemeDataset(Dataset):
-    def __init__(self, labels, img_dir, transform=None, target_transform=None):
-        self.img_labels = pd.read_csv(labels)
-        self.img_dir = img_dir
+    def __init__(self, data_dir, transform=None):
         self.transform = transform
-        self.target_transform = target_transform
+        self.img_path = []
+        for e in supported_ext:
+            path = glob.glob(os.path.join(data_dir, '*', f'*.{e}'))
+            for p in path:
+                if imghdr.what(p) == 'jpeg':
+                    self.img_path.append(p)
+        classes = set()
+        for path in self.img_path:
+            classes.add(os.path.basename(os.path.dirname(path)))
+        self.labels = {cls: i for i, cls in enumerate(sorted(list(classes)))}
+        self.resize = Resize(size = (640, 640))
 
     def __len__(self):
-        return len(self.img_labels)
-
+        return len(self.img_path)
+    
     def __getitem__(self, idx):
-        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
-        image = self.img_labels.iloc[idx, 1]
+        img = read_image(self.img_path[idx], ImageReadMode.RGB).float()
+        cls = os.path.basename(os.path.dirname(self.img_path[idx]))
+        label = self.labels[cls]
+
         if self.transform:
-            image = self.transform(image)
-        if self.target_transform:
-            label = self.target_transform(label)
-        return image, label
+            return self.transform(img), torch.tensor(label)
+
+        return self.resize(img), torch.tensor(label)

mACHINE-LEARNINGS/eval.py

@@ -0,0 +1,52 @@
+import warnings; warnings.filterwarnings('ignore')
+import torch
+from torch.utils.data import DataLoader
+from torchvision.transforms import Resize
+from torchvision.io import read_image, ImageReadMode
+from classes.MemeDataset import MemeDataset
+from models.DankCNN import DankCNN
+import argparse
+
+cls = {0: 'dank', 1: 'normie'}
+
+def eval(image):
+    """
+    SETUP
+    """
+    if torch.cuda.is_available():
+        device = torch.device('cuda')
+    else:
+        device = torch.device('cpu')
+
+    """
+    DATA LOADING
+        - Load all data: train, test, validation
+    """
+    image = read_image(image, ImageReadMode.RGB).float()
+    resize = Resize(size = (640, 640))
+    image = resize(image)
+
+    """
+    MODEL INITIALIZATION
+        - optimizer: Adam with weight decay as regularization technique
+        - loss function: binary cross entropy loss
+    """
+    model = DankCNN()
+    model.load_state_dict(torch.load('models/trained/model.pt'))
+    model = model.to(device)
+
+    model.eval()
+    with torch.no_grad():
+        image = image.unsqueeze(0)
+        output = model(image.to(device))
+        prediction = output.data
+        print('This meme is', cls[0 if prediction.item() < .5 else 1])
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-i', '--image', nargs=1, help='Image to evaluate', required=True)
+    args = parser.parse_args()
+    eval(
+        image=args.image[0]
+        )

mACHINE-LEARNINGS/models/DankCNN.py

@@ -4,16 +4,16 @@
 
 class DankCNN(nn.Module):
 
-    def __init__(self, dropout):
+    def __init__(self, dropout=False):
         super(DankCNN, self).__init__()
         self.conv1 = nn.Conv2d(in_channels=3, out_channels=6, kernel_size=4)
         self.pool = nn.MaxPool2d(kernel_size=3, stride=2)
         self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=4)
         self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
         self.conv3 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=4)
-        self.linear_layer1 = nn.Linear(21632, 128)
+        self.linear_layer1 = nn.Linear(758912, 128)
         self.linear_layer2 = nn.Linear(128, 96)
-        self.linear_layer3 = nn.Linear(96, 2)
+        self.linear_layer3 = nn.Linear(96, 1)
         self.dropout = nn.Dropout(p=.3)
         self.dropout_flag = dropout
 
@@ -33,4 +33,4 @@ def forward(self, x):
         x = F.relu(x)
         x = self.dropout(x) if self.dropout_flag else x
         x = self.linear_layer3(x)
-        return F.sigmoid(x)
+        return F.sigmoid(x)

mACHINE-LEARNINGS/train.py

@@ -1,15 +1,12 @@
-import numpy as np
-import pandas as pd
-import torch
 import warnings; warnings.filterwarnings('ignore')
-from classes.NumberDataset import NumberDataset
+import torch
+from torch.utils.data import DataLoader
+from classes.MemeDataset import MemeDataset
 from utils.performance_measure import precision_recall_f1
-from models.EvenNet import EvenNet
-import time
+from models.DankCNN import DankCNN
 import argparse
-import pickle
 
-def train(batch_size_train, lr, epochs, is_verbose, weight_decay):
+def execute(train_set_size, batch_size, lr, epochs, is_verbose, weight_decay):
     """
     HYPERPARAMETERS AND CONSTANTS
         - BATCH_SIZE_TRAIN: size of the batches for training phase
@@ -18,7 +15,9 @@ def train(batch_size_train, lr, epochs, is_verbose, weight_decay):
         - IS_VERBOSE: to avoid too much output
         - WEIGHT_DECAY: the weight decay for the regularization in Adam optimizer
     """
-    BATCH_SIZE_TRAIN = batch_size_train
+    TRAIN_SET_SIZE = train_set_size
+    TEST_SET_SIZE = 1 - train_set_size
+    BATCH_SIZE = batch_size
     LR = lr
     N_EPOCHS = epochs
     IS_VERBOSE = is_verbose
@@ -35,27 +34,93 @@ def train(batch_size_train, lr, epochs, is_verbose, weight_decay):
     DATA LOADING
         - Load all data: train, test, validation
     """
-
-    model = EvenNet(input_dim=batch_size_train)
+    dataset = MemeDataset('data/', transform=None)
+    train_size = int(TRAIN_SET_SIZE*len(dataset))
+    test_size = len(dataset) - train_size
+    train_set, test_set = torch.utils.data.random_split(dataset, [train_size, test_size])
+    dataloaders = {
+        "train": DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True),
+        "test": DataLoader(test_set, batch_size=BATCH_SIZE, shuffle=False)
+    }
+
     """
     MODEL INITIALIZATION
         - optimizer: Adam with weight decay as regularization technique
         - loss function: binary cross entropy loss
     """
+    model = DankCNN()
     model = model.to(device)
     optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)
-    loss_function = torch.nn.MSELoss()
+    loss_function = torch.nn.BCELoss()
+
+    """
+    TRAINING PHASE
+    """
+    for epoch in range(N_EPOCHS):
+        train_loss = 0
+        acc = 0
+
+        for batch_num, (image, label) in enumerate(dataloaders["train"]):
+            optimizer.zero_grad()
+            output = model(image.to(device))
+
+            label = label.unsqueeze(1)
+
+            loss = loss_function(output, label.float())
+            loss.backward()
+            optimizer.step()
+
+            train_loss += loss.item()
+
+            predictions = torch.where(output > .5, 1, 0)
+            acc += (label == predictions).sum()/len(label)
+
+            if IS_VERBOSE:
+                print('Training: Epoch %d - Batch %d/%d: Loss: %.4f' % 
+                (epoch+1, batch_num, len(dataloaders["train"]), train_loss / (batch_num + 1)))
+
+        print('EPOCH', epoch+1, 'ACCURACY:', (acc.item() / (train_size/BATCH_SIZE)))
+
+    """
+    TEST PHASE
+    """
+    test_loss = 0
+    acc = 0
+    model.eval()
+    with torch.no_grad():
+        for batch_num, (image, label) in enumerate(dataloaders["test"]):
+            output = model(image.to(device))
+
+            label = label.unsqueeze(1)
+            loss = loss_function(output, label.float())
+
+            test_loss += loss.item()
+
+            predictions = torch.where(output > .5, 1, 0)
+            acc += (label == predictions).sum()/len(label)
+
+            if IS_VERBOSE:
+                print('Evaluating: Batch %d/%d: Loss: %.4f' % 
+                    (batch_num, len(dataloaders["test"]), test_loss / (batch_num + 1)))
+
+        print('TEST ACCURACY:', (acc.item() / (test_size/BATCH_SIZE)))
+
+    print('Saving model...')
+    torch.save(model.state_dict(), 'models/trained/model.pt')
+    print('Saved 🌭')
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument('-ts', '--batchsizetrain', nargs=1, type=int, help='Size of the training batch', required=False)
+    parser.add_argument('-ts', '--trainsetsize', nargs=1, type=float, choices=[x/10 for x in range(0, 10)], help='Train set split size', required=True)
+    parser.add_argument('-bs', '--batchsize', nargs=1, type=int, help='Size of the training batch', required=False)
     parser.add_argument('-lr', '--learningrate', nargs=1, type=float, help='Learning rate', required=False)
     parser.add_argument('-e', '--epochs', nargs=1, type=int, help='Number of epochs', required=False)
     parser.add_argument('-v', '--verbose', nargs=1, type=bool, help='Verbose mode on/off', required=False)
     parser.add_argument('-wd', '--weightdecay', nargs=1, type=float, help='Weight decay (L2 regularization)', required=False)
     args = parser.parse_args()
-    train(
-        batch_size_train=args.batchsizetrain[0] if args.batchsizetrain else 1000,
+    execute(
+        train_set_size=args.trainsetsize[0] if args.trainsetsize[0]!=0 else 0.8,
+        batch_size=args.batchsize[0] if args.batchsize else 1000,
         lr=args.learningrate[0] if args.learningrate else 0.5,
         epochs=args.epochs[0] if args.epochs else 10,
         is_verbose=args.verbose if args.verbose else True,