2.EfficientNetB0.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import cv2\n",
    "import sys\n",
    "import random\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "from glob import glob\n",
    "from copy import deepcopy\n",
    "from skimage.feature import graycomatrix\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "import timm\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import albumentations as A\n",
    "import torch.optim as optim\n",
    "from ultralytics import YOLO\n",
    "import torch.nn.functional as F\n",
    "from torchvision import transforms\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from albumentations.pytorch.transforms import ToTensorV2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomModel(nn.Module):\n",
    "    def __init__(self, num_classes=1):\n",
    "        super().__init__()\n",
    "        self.effi = timm.create_model(\"efficientnet_b0\", pretrained=True)\n",
    "        self.effi.bn2 = nn.Conv2d(1280,64,kernel_size =1, stride =1)\n",
    "        self.effi.drop = nn.Identity()\n",
    "        self.effi.act = nn.Identity()\n",
    "        self.effi.global_pool = nn.Identity()\n",
    "        self.effi.classifier = nn.Flatten()\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(4096,1000),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Dropout(p=0.3),\n",
    "            nn.Linear(1000, 1)\n",
    "        )\n",
    "    \n",
    "    def forward(self, x):\n",
    "        # concat with entropy map\n",
    "        x = self.effi(x)\n",
    "        x = self.effi.global_pool(x)\n",
    "        x = self.effi.classifier(x)\n",
    "        x = self.fc(x)\n",
    "        x = F.sigmoid(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "  \n",
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, video_paths: list, labels: list=None, mode=\"train\", n=1, randomness=False):\n",
    "        assert mode in ['train', 'test', 'validation','val']\n",
    "        self.video_paths = video_paths\n",
    "        self.labels = labels\n",
    "        self.mode = mode\n",
    "        self.n = n\n",
    "        self.randomness = randomness\n",
    "        self.yolo = torch.hub.load('ultralytics/yolov5', 'yolov5s')\n",
    "        self.transform = A.Compose([A.Resize(480,480),\n",
    "                                    A.CenterCrop(256,256,p=1),\n",
    "                                    A.Normalize(0.5,0.5),\n",
    "                                    ToTensorV2()])\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.video_paths)\n",
    "    \n",
    "    def _get_co_occurrence_matrix(self,frames):\n",
    "        h, w = frames.shape\n",
    "\n",
    "        # 각 채널에 대한 gray scale로 변환\n",
    "        gray_images = 0.299 * frames[:, 0, :, :] + 0.587 * frames[:, 1, :, :] + 0.114 * frames[:, 2, :, :]\n",
    "        gray_images = gray_images.unsqueeze(1)  # (n, 1, h, w)\n",
    "\n",
    "        # 각 채널에 대한 필터 정의\n",
    "        filters_r = torch.eye(3).view(3, 1, 1, 1)\n",
    "        filters_g = torch.eye(3).view(3, 1, 1, 1)\n",
    "        filters_b = torch.eye(3).view(3, 1, 1, 1)\n",
    "\n",
    "        # 각 채널에 대한 co-occurrence matrix 계산\n",
    "        co_occurrence_matrices_r = F.conv2d(gray_images[:, [0]], filters_r, padding=1)\n",
    "        co_occurrence_matrices_g = F.conv2d(gray_images[:, [1]], filters_g, padding=1)\n",
    "        co_occurrence_matrices_b = F.conv2d(gray_images[:, [2]], filters_b, padding=1)\n",
    "\n",
    "        # 계산된 co-occurrence matrices 합치기\n",
    "        co_occurrence_matrices = torch.cat([co_occurrence_matrices_r, co_occurrence_matrices_g, co_occurrence_matrices_b], dim=1)\n",
    "\n",
    "        return co_occurrence_matrices\n",
    "    \n",
    "    def _get_video_frames(self, cap):\n",
    "        num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
    "        frames = []\n",
    "        \n",
    "        trial = 0\n",
    "        while len(frames) < self.n:\n",
    "            if self.randomness:\n",
    "                move_to = random.randint(1, num_frames-10*self.n)\n",
    "                cap.set(cv2.CAP_PROP_POS_FRAMES, move_to)\n",
    "            ret, frame = cap.read()\n",
    "            if ret: # 프레임 존재\n",
    "                results = self.yolo([frame])\n",
    "                df = results.pandas().xyxy[0]\n",
    "                df = df[df['name']=='person']\n",
    "                if len(df) >= 1:\n",
    "                    xmin, ymin, xmax, ymax, _,_,_ = df.iloc[0]\n",
    "                    xmin, ymin, xmax, ymax = int(xmin), int(ymin), int(xmax), int(ymax)\n",
    "                    frame = frame[ymin:ymax,xmin:xmax,:]\n",
    "                    frame = self.transform(image=frame)['image']\n",
    "                    frames.append(frame)\n",
    "                else: # 프레임이 존재하나 욜로 모델이 검출을 못함\n",
    "                    trial += 1\n",
    "                    move_to = random.randint(1, num_frames-10*self.n)\n",
    "                    cap.set(cv2.CAP_PROP_POS_FRAMES, move_to)\n",
    "                    if trial == 3:\n",
    "                        trial = 0\n",
    "                        ret, frame = cap.read()\n",
    "                        frame = self.crop(image=frame)['image']\n",
    "                        frames.append(frame)\n",
    "                    else:\n",
    "                        continue\n",
    "            else:\n",
    "                trial += 1\n",
    "                move_to = random.randint(1, num_frames-10*self.n)\n",
    "                cap.set(cv2.CAP_PROP_POS_FRAMES, move_to)\n",
    "                if trial == 3:\n",
    "                    trial = 0\n",
    "                    ret, frame = cap.read()\n",
    "                    frame = self.crop(image=frame)['image']\n",
    "                    frames.append(frame)\n",
    "                else:\n",
    "                    continue\n",
    "        frames = torch.stack(frames)\n",
    "        frames = frames.detach().clone()\n",
    "        return frames\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        video_path = self.video_paths[index]\n",
    "        # 랜덤 프레임 가져오기\n",
    "        cap = cv2.VideoCapture(video_path)\n",
    "        frames = self._get_video_frames(cap)\n",
    "        r_channel = frames[:, :, 2]\n",
    "        g_channel = frames[:, :, 1]\n",
    "        b_channel = frames[:, :, 0]\n",
    "        \n",
    "        distance = [1]\n",
    "        angles = [0, np.pi/4, np.pi/2, 3*np.pi/4]\n",
    "        red_matrix = graycomatrix(r_channel, distances=distance, angles=angles, symmetric=True, normed=True)\n",
    "        green_matrix = graycomatrix(g_channel, distances=distance, angles=angles, symmetric=True, normed=True)\n",
    "        blue_matrix = graycomatrix(b_channel, distances=distance, angles=angles, symmetric=True, normed=True)\n",
    "        combined_matrix = np.concatenate([red_matrix, green_matrix, blue_matrix], axis=-1)\n",
    "        frames = combined_matrix.squeeze(2)\n",
    "        cap.release()\n",
    "        if self.mode == 'test':\n",
    "            return frames\n",
    "        else:\n",
    "            return frames, self.labels[index]\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Trainer:\n",
    "    def __init__(self):\n",
    "        self.train_path = '/mnt/elice/dataset/train'\n",
    "        self.test_path = '/mnt/elice/dataset/test'\n",
    "        self.submission_csv = \"./sample_submission_v0.csv\"\n",
    "        self.save_path = \"./best_effi_model.pth\"\n",
    "        self.EPOCHS = 30\n",
    "        self.LR = 0.001\n",
    "        self.BATCH_SIZE=16\n",
    "        self.MAX_NORM = 5\n",
    "        self.NUM_WORKERS = 0\n",
    "        self.NUM_CLASSES = 1\n",
    "        self.DEVICE = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "        self.SEED = 777\n",
    "\n",
    "    def setup(self):\n",
    "        \n",
    "        seed_everything(self.SEED)\n",
    "        # 재현성 위해 sorting\n",
    "        train_fakes = sorted(glob(f\"{self.train_path}/fake/*\"))\n",
    "        train_reals = sorted(glob(f\"{self.train_path}/real/*\"))\n",
    "        self.submit = pd.read_csv(self.submission_csv)\n",
    "        x_test = [os.path.join(self.test_path, path) for path in self.submit[\"path\"].values]\n",
    "        # test_video_paths = sorted(glob(f\"{self.train_path}/*\"))\n",
    "        # fake이면 1 real이면 0으로 할당\n",
    "        train_video_paths = train_fakes + train_reals\n",
    "        labels = [1 for _ in range(len(train_fakes))] + [0 for _ in range(len(train_reals))]\n",
    "        x_train, x_val, y_train, y_val = train_test_split(\n",
    "            train_video_paths,\n",
    "            labels,\n",
    "            test_size=0.2,\n",
    "            random_state=self.SEED\n",
    "        )\n",
    "        \n",
    "        train_dataset = CustomDataset(video_paths=x_train, labels=y_train, mode=\"train\")\n",
    "        val_dataset = CustomDataset(video_paths=x_val, labels=y_val, mode=\"val\")\n",
    "        test_dataset = CustomDataset(video_paths=x_test, labels=None, mode=\"test\")\n",
    "        \n",
    "        self.train_dataloader = DataLoader(\n",
    "            dataset=train_dataset,\n",
    "            batch_size=self.BATCH_SIZE,\n",
    "            shuffle=True,\n",
    "            num_workers=self.NUM_WORKERS,\n",
    "            pin_memory=True\n",
    "        )\n",
    "        \n",
    "        self.val_dataloader = DataLoader(\n",
    "            dataset=val_dataset, \n",
    "            batch_size=self.BATCH_SIZE,\n",
    "            shuffle=False\n",
    "        )\n",
    "        \n",
    "        self.test_dataloader = DataLoader(\n",
    "            dataset=test_dataset, \n",
    "            batch_size=1,\n",
    "            shuffle=False\n",
    "        )\n",
    "        \n",
    "        self.model = CustomModel()\n",
    "        self.loss_fn = nn.BCELoss()\n",
    "        \n",
    "        \n",
    "    def train(self):\n",
    "        self.model.to(self.DEVICE)\n",
    "        optimizer = optim.AdamW(params=self.model.parameters(), lr=self.LR, weight_decay=1e-3)\n",
    "        lr_scheduler = optim.lr_scheduler.ReduceLROnPlateau(\n",
    "            optimizer=optimizer,\n",
    "            mode='max',\n",
    "            factor=0.5,\n",
    "            patience=3,\n",
    "            cooldown=5,\n",
    "            min_lr=1e-9,\n",
    "            threshold_mode='abs',\n",
    "        )\n",
    "        # AMP : loss scale을 위한 gradscaler\n",
    "        scaler = torch.cuda.amp.GradScaler()\n",
    "        \n",
    "        best_val_acc = 0\n",
    "        best_model = None\n",
    "        \n",
    "        for epoch in range(1, self.EPOCHS+1):\n",
    "            self.model.train()\n",
    "            train_losses = []\n",
    "            for idx,(imgs, labels) in tqdm(enumerate(self.train_dataloader)):\n",
    "                imgs = torch.squeeze(imgs, 1)\n",
    "                imgs = imgs.float().to(self.DEVICE) # (b,3,h,w)\n",
    "                labels = labels.float().to(self.DEVICE)\n",
    "                \n",
    "                optimizer.zero_grad()\n",
    "                output = self.model(imgs)\n",
    "                output = output.squeeze(-1)\n",
    "                loss = self.loss_fn(output, labels)\n",
    "                \n",
    "                loss.backward()\n",
    "                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.MAX_NORM)\n",
    "                optimizer.step()\n",
    "                \n",
    "                train_losses.append(loss.item())\n",
    "                \n",
    "\n",
    "            val_loss, val_acc = self._valid()\n",
    "            train_loss = np.mean(train_losses)\n",
    "            print(f\"EPOCH: {epoch}, TRAIN LOSS: {train_loss:.4f}, VAL LOSS: {val_loss:.4f}, VAL ACC: {val_acc:.4f}\")\n",
    "\n",
    "            if lr_scheduler is not None:\n",
    "                lr_scheduler.step(val_acc)\n",
    "\n",
    "            if best_val_acc <= val_acc:\n",
    "                best_val_acc = val_acc\n",
    "                best_model = deepcopy(self.model)\n",
    "                torch.save(self.model.state_dict(), self.save_path)\n",
    "                early_stop = 0\n",
    "            else:\n",
    "                early_stop += 1\n",
    "\n",
    "            if early_stop > 7:\n",
    "                break\n",
    "    \n",
    "    def _valid(self):\n",
    "        self.model.eval()\n",
    "        val_losses = []\n",
    "        val_accs = []\n",
    "        with torch.no_grad():\n",
    "            for imgs, labels in tqdm(self.val_dataloader):\n",
    "                imgs = torch.squeeze(imgs, 1)\n",
    "                imgs = imgs.float().to(self.DEVICE)\n",
    "                labels = labels.float().to(self.DEVICE)\n",
    "                \n",
    "                probs = self.model(imgs)\n",
    "                probs = probs.squeeze(-1)\n",
    "                loss = self.loss_fn(probs, labels)\n",
    "                probs = probs.cpu().detach().numpy()\n",
    "                labels = labels.cpu().detach().numpy()\n",
    "\n",
    "                preds = probs > 0.5\n",
    "                batch_acc = (labels == preds).mean()\n",
    "                val_accs.append(batch_acc)\n",
    "                val_losses.append(loss.item())\n",
    "        \n",
    "        return np.mean(val_losses), np.mean(val_accs)\n",
    "    \n",
    "    def test(self, threshold=0.5):\n",
    "        answer_lst = []\n",
    "        logit_lst = []\n",
    "        logit_df = deepcopy(self.submit)\n",
    "        # model load\n",
    "        model = CustomModel(self.NUM_CLASSES).to(self.DEVICE)\n",
    "        model.load_state_dict(torch.load(self.save_path))\n",
    "        model.eval()\n",
    "        with torch.no_grad():\n",
    "            for imgs in tqdm(self.test_dataloader):\n",
    "                cur_ans = []\n",
    "                for img in imgs:\n",
    "                    imgs = torch.squeeze(imgs, 1)\n",
    "                    imgs = imgs.float().to(self.DEVICE)\n",
    "\n",
    "                    probs = model(imgs)\n",
    "                    probs = probs.squeeze(-1)\n",
    "                    probs = probs.cpu().detach().numpy()\n",
    "                    logit_lst.append(probs[0])\n",
    "\n",
    "                    preds = 1 if probs > threshold else 0\n",
    "                    cur_ans.append(preds)\n",
    "                cnt_1 = cur_ans.count(1)\n",
    "                cnt_0 = cur_ans.count(0)\n",
    "                ans = 1 if cnt_1 > cnt_0 else 0\n",
    "                answer_lst.append(ans)\n",
    "\n",
    "            self.submit[\"label\"] = answer_lst\n",
    "            self.submit[\"label\"] = self.submit[\"label\"].apply(lambda x: \"fake\" if x else \"real\")\n",
    "            self.submit.to_csv(\"sample_submissionv2.csv\", index=False)\n",
    "            \n",
    "            logit_df['label'] = logit_lst\n",
    "            logit_df.to_csv(\"sample_submission_logitv2.csv\", index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer()\n",
    "trainer.setup()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.test()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}