ELA_ResNeXt_Classifier.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "059b39e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.optim as optim\n",
    "from torch.nn import CrossEntropyLoss\n",
    "import torch.nn.functional as F\n",
    "import torch.nn as nn\n",
    "from torchvision.datasets.folder import DatasetFolder\n",
    "import torchvision.transforms as transforms \n",
    "from torchvision import models\n",
    "from skimage.io import imread\n",
    "from sklearn import metrics\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import pandas as pd\n",
    "import cv2\n",
    "import numpy as np\n",
    "import random\n",
    "from PIL import Image\n",
    "import sklearn\n",
    "from flopth import flopth\n",
    "import os\n",
    "from sklearn.model_selection import train_test_split\n",
    "from PIL import Image, ImageChops\n",
    "from torchvision.models.resnet import ResNeXt50_32X4D_Weights\n",
    "from data import data_preprocess"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28aa9cc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 20\n",
    "lr = 0.01\n",
    "momentum = 0.5\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda:0\")\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "    \n",
    "DATA_DIR = './'\n",
    "    \n",
    "class MyDataset(Dataset):\n",
    "    def __init__(self, root_dir, test_size=0.2, transform=None, test = False):\n",
    "        \n",
    "        #Collect all samples from our dataset\n",
    "        #Note that our dataset contains 5000 true and 5000 photoshopped images\n",
    "        self.root_dir = root_dir\n",
    "        self.transform = transform\n",
    "        \n",
    "        self.class_folders = ['originals', 'photoshops']\n",
    "        self.class_folders_ela = ['originals_ela', 'photoshops_ela']\n",
    "        self.samples = []\n",
    "\n",
    "        for class_folder, class_folder_ela in zip(self.class_folders, self.class_folders_ela):\n",
    "          \n",
    "            class_path = os.path.join(self.root_dir, class_folder)\n",
    "            class_path_ela = os.path.join(self.root_dir, class_folder_ela)\n",
    "            class_label = self.class_folders.index(class_folder)\n",
    "            filenames = sorted(os.listdir(class_path))\n",
    "            filenames_ela = sorted(os.listdir(class_path_ela))\n",
    "          \n",
    "            for file_name, file_name_ela in zip(filenames, filenames_ela):\n",
    "                file_path = os.path.join(class_path, file_name)\n",
    "                file_path_ela = os.path.join(class_path_ela, file_name_ela)\n",
    "                self.samples.append((file_path, file_path_ela, class_label))\n",
    "        \n",
    "        # randomly split into train and test sets\n",
    "        train_samples, test_samples = train_test_split(self.samples, test_size=test_size)\n",
    "        \n",
    "        #Subsample data as needed for experiments\n",
    "        #train_samples = train_samples[:4000]\n",
    "        #test_samples = test_samples[:1000]\n",
    "        \n",
    "        #Label data\n",
    "        if test:\n",
    "            self.samples = test_samples\n",
    "            self.train = False\n",
    "        else:\n",
    "            self.samples = train_samples\n",
    "            self.train = True\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.samples)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        file_path, file_path_ela, class_label = self.samples[index]\n",
    "        image = Image.open(file_path).convert('RGB')\n",
    "        image_ela = Image.open(file_path_ela).convert('RGB')\n",
    "\n",
    "        if self.transform is not None:\n",
    "            image = self.transform(image)\n",
    "            image_ela = self.transform(image_ela)\n",
    "\n",
    "        return image, image_ela, class_label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeb00658",
   "metadata": {},
   "outputs": [],
   "source": [
    "#transform and split data\n",
    "my_transform = transforms.Compose([\n",
    "    transforms.Resize((224, 224)), \n",
    "    transforms.ToTensor(),\n",
    "    transforms.RandomVerticalFlip(0.5),\n",
    "    transforms.RandomHorizontalFlip(0.5),\n",
    "    transforms.RandomRotation(15),\n",
    "    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n",
    "])\n",
    "\n",
    "dataset = MyDataset(DATA_DIR, test_size=0.2, transform=my_transform, test = False)\n",
    "train_loader = DataLoader(dataset, batch_size=32, shuffle=True)\n",
    "\n",
    "test_dataset = MyDataset(DATA_DIR, test_size=0.2, transform=my_transform, test = True)\n",
    "test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)\n",
    "\n",
    "print(len(train_loader.dataset))\n",
    "print(len(test_loader.dataset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c388a65",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ELRes(nn.Module):\n",
    "    def __init__(self, num_classes):\n",
    "        super(CustomCNN, self).__init__()\n",
    "        self.hidden_size = 4096\n",
    "        self.modelELA = models.resnext50_32x4d(weights=ResNeXt50_32X4D_Weights.DEFAULT)\n",
    "        self.modelRGB = models.resnext50_32x4d(weights=ResNeXt50_32X4D_Weights.DEFAULT)\n",
    "        self.modelELA.fc = nn.Identity()\n",
    "        self.modelRGB.fc = nn.Identity()\n",
    "        self.fc = nn.Linear(self.hidden_size, num_classes)\n",
    "        \n",
    "    def forward(self, x1, x2):\n",
    "        ela_features = self.modelELA(x1)\n",
    "        rgb_features = self.modelRGB(x2)\n",
    "        combined_features = torch.cat([ela_features, rgb_features], dim=1)\n",
    "        self.hidden_size = combined_features.shape[0]\n",
    "        output = self.fc(combined_features)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfe2564a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = ELRes(num_classes = 2).to(device)\n",
    "criterion = CrossEntropyLoss()\n",
    "optimizer = optim.SGD(model.parameters(),lr=lr,momentum=momentum)\n",
    "#scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67e0066c",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "for epoch_num in range(0, epochs):\n",
    "    model.train()\n",
    "    epoch_total_loss = 0\n",
    "    train_predictions = []\n",
    "    train_labels = []\n",
    "    for batch_num, (inp1, inp2, target) in enumerate(train_loader):\n",
    "        if batch_num % 10 == 0:\n",
    "            print(\"EPOCH\", epoch_num, \"Batch number\", batch_num)\n",
    "        train_labels+=target\n",
    "        optimizer.zero_grad()\n",
    "        output = model(inp1.to(device), inp2.to(device))\n",
    "        batch_loss = criterion(output,target.to(device))\n",
    "        _, prediction = torch.max(output, dim=1)\n",
    "        train_predictions += prediction.detach().tolist()\n",
    "        epoch_total_loss += batch_loss.item()\n",
    "        batch_loss.backward()\n",
    "        optimizer.step()\n",
    "    avrg_loss = epoch_total_loss / dataset.__len__()\n",
    "    train_accuracy = metrics.accuracy_score(train_labels, train_predictions)\n",
    "    print(\"Train Accuracy = %0.2f\" % (train_accuracy))\n",
    "    print(\"Epoch %d - loss=%0.4f\" % (epoch_num, avrg_loss))\n",
    "    #scheduler.step()\n",
    "\n",
    "\n",
    "    model.eval()\n",
    "    labels = []\n",
    "    predictions = []\n",
    "    for batch_num, (inp1, inp2, target) in enumerate(test_loader):\n",
    "        labels+=target\n",
    "        batch_prediction = model(inp1.to(device), inp2.to(device))\n",
    "        _, batch_prediction = torch.max(batch_prediction, dim=1)\n",
    "        predictions += batch_prediction.detach().tolist()\n",
    "    accuracy = metrics.accuracy_score(labels, predictions)\n",
    "    print(\"Test Accuracy = %0.2f\" % (accuracy))\n",
    "    confusion = metrics.confusion_matrix(labels, predictions)\n",
    "\n",
    "    try:\n",
    "        print(confusion)\n",
    "        f1_score = sklearn.metrics.f1_score(labels, predictions)\n",
    "        print(f1_score)\n",
    "        recall = sklearn.metrics.recall_score(labels, predictions)\n",
    "    except:\n",
    "        pass"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}