surver.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pred Distance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd \n",
    "import glob\n",
    "import cv2 as cv\n",
    "import random\n",
    "import os\n",
    "import numpy as np\n",
    "import random\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import albumentations as A\n",
    "from albumentations.pytorch.transforms import ToTensorV2\n",
    "from tqdm.auto import tqdm\n",
    "import timm\n",
    "\n",
    "\n",
    "review_img = cv.imread('1.png')\n",
    "product_img = cv.imread('2.png')\n",
    "\n",
    "def pred_distance(review_img_path, product_img_path):\n",
    "    df = pd.DataFrame(columns=['review_img_path','product_img_path', 'label'])\n",
    "    df['review_img_path'] = [review_img_path]\n",
    "    df['product_img_path'] = [product_img_path]\n",
    "    df['label'] = [0]\n",
    "    \n",
    "    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    "    \n",
    "    CFG = {\n",
    "        'IMG_SIZE':224,\n",
    "        'EPOCHS':1,\n",
    "        'LEARNING_RATE':3e-4,\n",
    "        # 'LEARNING_RATE':10,\n",
    "        'BATCH_SIZE':1,\n",
    "        'SEED':41\n",
    "    }\n",
    "    \n",
    "    class SiameseNetworkDataset(Dataset):\n",
    "        def __init__(self,review_img_path,product_img_path,label,transform=None):\n",
    "            self.review_img_path = review_img_path\n",
    "            self.product_img_path = product_img_path\n",
    "            self.label = label\n",
    "            self.transform = transform\n",
    "\n",
    "        def __getitem__(self,index):\n",
    "            # review_img = cv.imread(self.review_img_path[index])\n",
    "            # product_img = cv.imread(self.product_img_path[index])\n",
    "            review_img = self.review_img_path[index]\n",
    "            product_img = self.product_img_path[index]\n",
    "            review_img = cv.resize(review_img, (CFG['IMG_SIZE'], CFG['IMG_SIZE']))\n",
    "            product_img = cv.resize(product_img, (CFG['IMG_SIZE'], CFG['IMG_SIZE']))\n",
    "\n",
    "            if self.transform is not None:\n",
    "                review_img  = self.transform(image=review_img)['image']\n",
    "                product_img  = self.transform(image=product_img)['image']\n",
    "\n",
    "            return review_img, product_img, self.label[index]\n",
    "\n",
    "        def __len__(self):\n",
    "            return len(self.review_img_path)\n",
    "        \n",
    "    train_transform = A.Compose([\n",
    "                            A.Resize(CFG['IMG_SIZE'],CFG['IMG_SIZE']),\n",
    "                            ToTensorV2()\n",
    "                            ])\n",
    "    \n",
    "    val_dataset = SiameseNetworkDataset(df[\"review_img_path\"].values, df[\"product_img_path\"].values, df[\"label\"].values, train_transform)\n",
    "    val_loader = DataLoader(val_dataset, batch_size = CFG['BATCH_SIZE'], shuffle=False, num_workers=0)\n",
    "\n",
    "    class BaseModel(nn.Module):\n",
    "        def __init__(self):\n",
    "            super(BaseModel, self).__init__()\n",
    "            self.backbone = timm.create_model('efficientnet_b0', pretrained=False)\n",
    "            self.classifier = nn.Linear(1000, 50)\n",
    "            self.dropout = nn.Dropout(0.1)\n",
    "            self.ReLU = nn.ReLU(inplace=False)\n",
    "\n",
    "        def forward(self, x, y):\n",
    "            x = self.backbone(x)\n",
    "            x = self.classifier(x)\n",
    "\n",
    "            y = self.backbone(y)\n",
    "            y = self.classifier(y)\n",
    "\n",
    "            z = F.pairwise_distance(x, y, keepdim = True)\n",
    "            return z\n",
    "        \n",
    "    def validation(model, val_loader, device):\n",
    "        model.eval()\n",
    "        pred_list = []\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            for review_img, product_img, labels in iter(val_loader):\n",
    "                review_img = review_img.float().to(device)\n",
    "                product_img = product_img.float().to(device)\n",
    "    \n",
    "                \n",
    "                pred = model(review_img, product_img)\n",
    "                pred = pred.detach().cpu().numpy().tolist()\n",
    "                pred_list += pred\n",
    "            \n",
    "        return pred_list \n",
    "        \n",
    "    def train(model,  val_loader,  device):\n",
    "        model = model.to(device)\n",
    "\n",
    "        model.train()\n",
    "        prediction = validation(model, val_loader, device)\n",
    "\n",
    "        return prediction\n",
    "        \n",
    "    \n",
    "    model = BaseModel()\n",
    "    model.load_state_dict(torch.load('./distance_EffNetBase_E_Contra.pt'))\n",
    "    model.eval()\n",
    "\n",
    "    prediction = train(model,  val_loader, device)\n",
    "    \n",
    "    return prediction[0][0]\n",
    "    \n",
    "# print(pred_distance(review_img, product_img))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pose Estimate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import mediapipe as mp\n",
    "from mediapipe.tasks import python\n",
    "from mediapipe.tasks.python import vision\n",
    "import numpy as np\n",
    "import mediapipe as mp\n",
    "from mediapipe.tasks import python\n",
    "from mediapipe.tasks.python import vision\n",
    "from tqdm.notebook import tqdm\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import random\n",
    "import os\n",
    "from sklearn.model_selection import train_test_split\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "import torch.nn.functional as F\n",
    "\n",
    "\n",
    "base_options = python.BaseOptions(model_asset_path='pose_landmarker.task')\n",
    "options = vision.PoseLandmarkerOptions(\n",
    "        base_options=base_options,\n",
    "        num_poses = 22,\n",
    "        output_segmentation_masks=False)\n",
    "\n",
    "detector = vision.PoseLandmarker.create_from_options(options)\n",
    "\n",
    "def Pose_Estimation(img_path):\n",
    "    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    "    CFG = {\n",
    "        'EPOCHS':1,\n",
    "        'LEARNING_RATE':3e-8,\n",
    "        # 'LEARNING_RATE':10,\n",
    "        'BATCH_SIZE':1,\n",
    "        'SEED':41\n",
    "        }\n",
    "    \n",
    "    save_x = []\n",
    "    save_y = []\n",
    "    save_z = []\n",
    "    save_presence = []\n",
    "    \n",
    "        \n",
    "    img = mp.Image.create_from_file(img_path)\n",
    "    pose_landmarks_list = detector.detect(img).pose_landmarks\n",
    "        \n",
    "    if not pose_landmarks_list:\n",
    "        return False\n",
    "            \n",
    "    save_x.append([i.x for i in pose_landmarks_list[0][11:33]])\n",
    "    save_y.append([i.y for i in pose_landmarks_list[0][11:33]])\n",
    "    save_z.append([i.z for i in pose_landmarks_list[0][11:33]])\n",
    "    save_presence.append([i.presence for i in pose_landmarks_list[0][11:33]])\n",
    "    \n",
    "    df = pd.DataFrame(columns=['img_path','label'])\n",
    "    \n",
    "    df['img_path'] = None\n",
    "    df['label'] = 0\n",
    "    df['landmark_x'] = save_x\n",
    "    df['landmark_y'] = save_y\n",
    "    df['landmark_z'] = save_z\n",
    "    df['landmark_presence'] = save_presence\n",
    "    \n",
    "    class CustomDataset(Dataset):\n",
    "        def __init__(self, img_path ,landmark_x, landmark_y,\tlandmark_z,\tlandmark_presence, label):\n",
    "            self.img_path = img_path\n",
    "            self.landmark_x = landmark_x\n",
    "            self.landmark_y = landmark_y\n",
    "            self.landmark_z= landmark_z\n",
    "            self.landmark_presence = landmark_presence\n",
    "            self.label = label\n",
    "\n",
    "        def __getitem__(self,index):\n",
    "\n",
    "            result = np.concatenate((self.landmark_x[index] , self.landmark_y[index] , self.landmark_z[index] , self.landmark_presence[index]), axis=0)\n",
    "            return result, self.label[index]\n",
    "\n",
    "        def __len__(self):\n",
    "            return len(self.landmark_x )\n",
    "    \n",
    "    val_dataset = CustomDataset(df[\"img_path\"].values, df[\"landmark_x\"].values, df[\"landmark_y\"].values, df[\"landmark_z\"].values,df[\"landmark_presence\"].values, df[\"label\"].values)\n",
    "    val_loader = DataLoader(val_dataset, batch_size = CFG['BATCH_SIZE'], shuffle=False, num_workers=0)\n",
    "    \n",
    "    class BaseModel(nn.Module):\n",
    "        def __init__(self):\n",
    "            super(BaseModel, self).__init__()\n",
    "            self.classifier1 = nn.Linear(88, 20)\n",
    "            # self.classifier1 = nn.Linear(22, 2)\n",
    "            self.ReLU = nn.ReLU(inplace=True)\n",
    "            self.classifier2 = nn.Linear(20, 2)\n",
    "\n",
    "\n",
    "        def forward(self, x):\n",
    "            x = self.classifier1(x)\n",
    "            x = self.ReLU(x)\n",
    "            x = self.classifier2(x)\n",
    "\n",
    "            return F.log_softmax(x, dim=1)\n",
    "        \n",
    "    def validation(model, criterion, val_loader, device):\n",
    "        model.eval()\n",
    "        preds= []\n",
    "        with torch.no_grad():\n",
    "            for landmark_list, labels in iter(val_loader):\n",
    "                landmark_list = landmark_list.float().to(device)\n",
    "                pred = model(landmark_list)\n",
    "                preds += pred.detach().argmax(1).cpu().numpy().tolist()\n",
    "        return preds\n",
    "\n",
    "\n",
    "    def train(model, optimizer, train_loader, val_loader, scheduler, device):\n",
    "        model = model.to(device)\n",
    "        criterion = nn.NLLLoss(weight=torch.tensor([0.01, 0.99]), reduction=\"sum\").to(device)\n",
    "        best_model = None\n",
    "        for epoch in range(0, CFG['EPOCHS']):\n",
    "            model.train()\n",
    "            label = validation(model, criterion, val_loader, device)\n",
    "        return best_model , label\n",
    "    \n",
    "    model = BaseModel()\n",
    "    model.load_state_dict(torch.load('./Pose_Estimate_new.pt'))\n",
    "    model.eval()\n",
    "    \n",
    "    optimizer = torch.optim.Adam(params = model.parameters(), lr = CFG[\"LEARNING_RATE\"])\n",
    "    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', factor=0.5, patience=2, threshold_mode='abs', min_lr=1e-8, verbose=True)\n",
    "    infer_model, label = train(model, optimizer, None, val_loader, scheduler, device)\n",
    "    \n",
    "    if label[0] == 0:\n",
    "        return True\n",
    "    else:\n",
    "        return False\n",
    "\n",
    "\n",
    "print(Pose_Estimation('./16.jpg'))\n",
    "    \n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# surver"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running on local URL:  http://127.0.0.1:7860\n",
      "Running on public URL: https://ddb64f851ac3f61675.gradio.live\n",
      "\n",
      "This share link expires in 72 hours. For free permanent hosting and GPU upgrades, run `gradio deploy` from Terminal to deploy to Spaces (https://huggingface.co/spaces)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"https://ddb64f851ac3f61675.gradio.live\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": []
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading model...\n"
     ]
    }
   ],
   "source": [
    "import gradio as gr\n",
    "import mask_image\n",
    "\n",
    "\n",
    "def image_classifier(review_img, product_img):\n",
    "    width = 30\n",
    "    height = 50\n",
    "    \n",
    "    if Pose_Estimation(review_img):\n",
    "        review_img = cv.imread(review_img)\n",
    "        product_img = cv.imread(product_img)\n",
    "        \n",
    "        img_list  = mask_image.mask_image([review_img, product_img])\n",
    "        preprocessed_review = img_list[0]\n",
    "        preprocessed_product = img_list[1] \n",
    "        \n",
    "        if (preprocessed_review is None) or (preprocessed_product is None):\n",
    "            return str(\"이미지를 분류하는데 실패하였습니다\"), False\n",
    "        \n",
    "        distance = pred_distance(preprocessed_review, preprocessed_product)\n",
    "        \n",
    "        if distance > 0.45:\n",
    "            return str(\"구매하신 제품이 아닙니다.\\n(distance :\" + str(distance) + \")\"), False\n",
    "        else:\n",
    "            return str(\"적립금이 지급되었습니다.\\n(distance :\" + str(distance) + \")\"), True\n",
    "    else:\n",
    "        return str(\"전신사진을 넣어주세요.\"), False\n",
    "    # return str(Pose_Estimation(review_img))\n",
    "\n",
    "\n",
    "product_1 = [None, \"./19.jpg\"]\n",
    "product_2 = [None, \"./20.jpg\"]\n",
    "product_3 = [None, \"./21.jpg\"]\n",
    "product_4 = [None, \"./18.jpg\"]\n",
    "product_5 = [None, \"./pro1.jpg\"]\n",
    "product_6 = [None, \"./pro2.jpg\"]\n",
    "\n",
    "\n",
    "\n",
    "image_input1 = gr.Image(type =\"filepath\",label= \"리뷰 이미지 올리세요\")\n",
    "image_input2 = gr.Image(type =\"filepath\",label= \"구매 제품을 선택하세요\")\n",
    "\n",
    "\n",
    "    \n",
    "\n",
    "demo = gr.Interface(\n",
    "                    fn=image_classifier,\n",
    "                    inputs= [image_input1,image_input2],\n",
    "                    outputs=\"text\",\n",
    "                    examples = [product_1,product_2,product_3,product_4,product_5,product_6]\n",
    "                    )\n",
    "demo.launch(share=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "SAM",
   "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.9.0"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}