segmentation/remove_background_app.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 배경 제거 Application\n",
    "\n",
    "Colab 환경에서 배경 제거 애플리케이션을 만들어봅시다. 애플리케이션 사용자의 유스케이스는 아래와 같습니다.\n",
    "\n",
    "- 사용자는 이미지 파일을 업로드할 수 있다.\n",
    "- 사용자는 이미지에서 원하는 객체 클릭한다.\n",
    "- 사용자는 배경 제거 이미지의 결과를 확인하고 다운로드 받을 수 있다."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 패키지 및 예제 데이터 다운로드하기\n",
    "python package들을 설치합니다. 예제로 사용할 이미지들도 다운로드 받습니다. Colab에서 실행하지 않는 경우 이 셀은 실행하지 않습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://raw.githubusercontent.com/mrsyee/dl_apps/main/segmentation/requirements-colab.txt\n",
    "!pip install -r requirements-colab.txt\n",
    "\n",
    "!mkdir examples\n",
    "!cd examples && wget https://github.com/mrsyee/dl_apps/raw/main/segmentation/examples/dog.jpg\n",
    "!cd examples && wget https://github.com/mrsyee/dl_apps/raw/main/segmentation/examples/mannequin.jpg"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 패키지 불러오기"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import urllib\n",
    "from typing import Tuple\n",
    "\n",
    "import cv2\n",
    "import gradio as gr\n",
    "import numpy as np\n",
    "import torch\n",
    "from PIL import Image\n",
    "from segment_anything import SamPredictor, sam_model_registry"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 애플리케이션 UI 구현하기"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with gr.Blocks() as app:\n",
    "    gr.Markdown(\"# Interactive Remove Background from Image\")\n",
    "    with gr.Row():\n",
    "        coord_x = gr.Number(label=\"Mouse coords x\")\n",
    "        coord_y = gr.Number(label=\"Mouse coords y\")\n",
    "\n",
    "    with gr.Row():\n",
    "        input_img = gr.Image(label=\"Input image\", height=600)\n",
    "        output_img = gr.Image(label=\"Output image\", height=600)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "app.launch(inline=False, share=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "app.close()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 마우스 클릭 이벤트 구현하기"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_coords(evt: gr.SelectData):\n",
    "    return evt.index[0], evt.index[1]\n",
    "\n",
    "\n",
    "with gr.Blocks() as app:\n",
    "    gr.Markdown(\"# Interactive Remove Background from Image\")\n",
    "    with gr.Row():\n",
    "        coord_x = gr.Number(label=\"Mouse coords x\")\n",
    "        coord_y = gr.Number(label=\"Mouse coords y\")\n",
    "\n",
    "    with gr.Row():\n",
    "        input_img = gr.Image(label=\"Input image\", height=600)\n",
    "        output_img = gr.Image(label=\"Output image\", height=600)\n",
    "\n",
    "    input_img.select(get_coords, None, [coord_x, coord_y])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "app.launch(inline=False, share=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "app.close()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SAM 추론기 클래스 구현하기"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CHECKPOINT_PATH = os.path.join(\"checkpoint\")\n",
    "CHECKPOINT_NAME = \"sam_vit_h_4b8939.pth\"\n",
    "CHECKPOINT_URL = \"https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth\"\n",
    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SAMInferencer:\n",
    "    def __init__(\n",
    "        self,\n",
    "        checkpoint_path: str,\n",
    "        checkpoint_name: str,\n",
    "        checkpoint_url: str,\n",
    "        model_type: str,\n",
    "        device: torch.device,\n",
    "    ):\n",
    "        print(\"[INFO] Initailize inferencer\")\n",
    "        if not os.path.exists(checkpoint_path):\n",
    "            os.makedirs(checkpoint_path, exist_ok=True)\n",
    "        checkpoint = os.path.join(checkpoint_path, checkpoint_name)\n",
    "        if not os.path.exists(checkpoint):\n",
    "            urllib.request.urlretrieve(checkpoint_url, checkpoint)\n",
    "        sam = sam_model_registry[model_type](checkpoint=checkpoint).to(device)\n",
    "        self.predictor = SamPredictor(sam)\n",
    "\n",
    "    def inference(\n",
    "        self,\n",
    "        image: np.ndarray,\n",
    "        point_coords: np.ndarray,\n",
    "        points_labels: np.ndarray,\n",
    "    ) -> np.ndarray:\n",
    "        self.predictor.set_image(image)\n",
    "        masks, scores, _ = self.predictor.predict(point_coords, points_labels)\n",
    "        mask, _ = self.select_mask(masks, scores)\n",
    "        return mask\n",
    "\n",
    "    def select_mask(\n",
    "        self, masks: np.ndarray, scores: np.ndarray\n",
    "    ) -> Tuple[np.ndarray, np.ndarray]:\n",
    "        # Determine if we should return the multiclick mask or not from the number of points.\n",
    "        # The reweighting is used to avoid control flow.\n",
    "        # Reference: https://github.com/facebookresearch/segment-anything/blob/6fdee8f2727f4506cfbbe553e23b895e27956588/segment_anything/utils/onnx.py#L92-L105\n",
    "        score_reweight = np.array([-1000] + [0] * 2)\n",
    "        score = scores + score_reweight\n",
    "        best_idx = np.argmax(score)\n",
    "        selected_mask = np.expand_dims(masks[best_idx, :, :], axis=-1)\n",
    "        selected_score = np.expand_dims(scores[best_idx], axis=0)\n",
    "        return selected_mask, selected_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inferencer = SAMInferencer(\n",
    "    CHECKPOINT_PATH, CHECKPOINT_NAME, CHECKPOINT_URL, \"vit_h\", DEVICE\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 추론 및 배경 제거 후처리 구현하기"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_object(image: np.ndarray, point_x: int, point_y: int):\n",
    "    point_coords = np.array([[point_x, point_y]])\n",
    "    point_label = np.array([1])\n",
    "\n",
    "    # Get mask\n",
    "    mask = inferencer.inference(image, point_coords, point_label)\n",
    "\n",
    "    # Extract object and remove background\n",
    "    # Postprocess mask\n",
    "    mask = (mask > 0).astype(np.uint8)\n",
    "\n",
    "    # Remove background\n",
    "    result_image = cv2.bitwise_and(image, image, mask=mask)\n",
    "\n",
    "    # Convert to rgba channel\n",
    "    bgr_channel = result_image[..., :3]  # BGR 채널 분리\n",
    "    alpha_channel = np.where(bgr_channel[..., 0] == 0, 0, 255).astype(np.uint8)\n",
    "    result_image = np.dstack((bgr_channel, alpha_channel))  # BGRA 이미지 생성\n",
    "\n",
    "    return result_image\n",
    "\n",
    "\n",
    "def extract_object_by_event(image: np.ndarray, evt: gr.SelectData):\n",
    "    click_x, click_y = evt.index\n",
    "\n",
    "    return extract_object(image, click_x, click_y)\n",
    "\n",
    "\n",
    "def get_coords(evt: gr.SelectData):\n",
    "    return evt.index[0], evt.index[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with gr.Blocks() as app:\n",
    "    gr.Markdown(\"# Interactive Remove Background from Image\")\n",
    "    with gr.Row():\n",
    "        coord_x = gr.Number(label=\"Mouse coords x\")\n",
    "        coord_y = gr.Number(label=\"Mouse coords y\")\n",
    "\n",
    "    with gr.Row():\n",
    "        input_img = gr.Image(label=\"Input image\", height=600)\n",
    "        output_img = gr.Image(label=\"Output image\", height=600)\n",
    "\n",
    "    input_img.select(extract_object_by_event, [input_img], [output_img])\n",
    "    input_img.select(get_coords, None, [coord_x, coord_y])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "app.launch(inline=False, share=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "app.close()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 최종 App 구현하기"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Implement inferencer\n",
    "CHECKPOINT_PATH = os.path.join(\"checkpoint\")\n",
    "CHECKPOINT_NAME = \"sam_vit_h_4b8939.pth\"\n",
    "CHECKPOINT_URL = \"https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth\"\n",
    "MODEL_TYPE = \"default\"\n",
    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "\n",
    "class SAMInferencer:\n",
    "    def __init__(\n",
    "        self,\n",
    "        checkpoint_path: str,\n",
    "        checkpoint_name: str,\n",
    "        checkpoint_url: str,\n",
    "        model_type: str,\n",
    "        device: torch.device,\n",
    "    ):\n",
    "        print(\"[INFO] Initailize inferencer\")\n",
    "        if not os.path.exists(checkpoint_path):\n",
    "            os.makedirs(checkpoint_path, exist_ok=True)\n",
    "        checkpoint = os.path.join(checkpoint_path, checkpoint_name)\n",
    "        if not os.path.exists(checkpoint):\n",
    "            urllib.request.urlretrieve(checkpoint_url, checkpoint)\n",
    "        sam = sam_model_registry[model_type](checkpoint=checkpoint).to(device)\n",
    "        self.predictor = SamPredictor(sam)\n",
    "\n",
    "    def inference(\n",
    "        self,\n",
    "        image: np.ndarray,\n",
    "        point_coords: np.ndarray,\n",
    "        points_labels: np.ndarray,\n",
    "    ) -> np.ndarray:\n",
    "        self.predictor.set_image(image)\n",
    "        masks, scores, _ = self.predictor.predict(point_coords, points_labels)\n",
    "        mask, _ = self.select_mask(masks, scores)\n",
    "        return mask\n",
    "\n",
    "    def select_mask(\n",
    "        self, masks: np.ndarray, scores: np.ndarray\n",
    "    ) -> Tuple[np.ndarray, np.ndarray]:\n",
    "        # Determine if we should return the multiclick mask or not from the number of points.\n",
    "        # The reweighting is used to avoid control flow.\n",
    "        # Reference: https://github.com/facebookresearch/segment-anything/blob/6fdee8f2727f4506cfbbe553e23b895e27956588/segment_anything/utils/onnx.py#L92-L105\n",
    "        score_reweight = np.array([-1000] + [0] * 2)\n",
    "        score = scores + score_reweight\n",
    "        best_idx = np.argmax(score)\n",
    "        selected_mask = np.expand_dims(masks[best_idx, :, :], axis=-1)\n",
    "        selected_score = np.expand_dims(scores[best_idx], axis=0)\n",
    "        return selected_mask, selected_score\n",
    "\n",
    "\n",
    "inferencer = SAMInferencer(\n",
    "    CHECKPOINT_PATH, CHECKPOINT_NAME, CHECKPOINT_URL, MODEL_TYPE, DEVICE\n",
    ")\n",
    "\n",
    "# Implement event function\n",
    "def extract_object(image: np.ndarray, point_x: int, point_y: int):\n",
    "    point_coords = np.array([[point_x, point_y], [0, 0]])\n",
    "    point_label = np.array([1, -1])\n",
    "\n",
    "    # Get mask\n",
    "    mask = inferencer.inference(image, point_coords, point_label)\n",
    "\n",
    "    # Extract object and remove background\n",
    "    # Postprocess mask\n",
    "    mask = (mask > 0).astype(np.uint8)\n",
    "\n",
    "    # Remove background\n",
    "    result_image = cv2.bitwise_and(image, image, mask=mask)\n",
    "\n",
    "    # Convert to rgba channel\n",
    "    bgr_channel = result_image[..., :3]  # BGR 채널 분리\n",
    "    alpha_channel = np.where(bgr_channel[..., 0] == 0, 0, 255).astype(np.uint8)\n",
    "    result_image = np.dstack((bgr_channel, alpha_channel))  # BGRA 이미지 생성\n",
    "\n",
    "    return result_image\n",
    "\n",
    "\n",
    "def extract_object_by_event(image: np.ndarray, evt: gr.SelectData):\n",
    "    click_x, click_y = evt.index\n",
    "\n",
    "    return extract_object(image, click_x, click_y)\n",
    "\n",
    "\n",
    "def get_coords(evt: gr.SelectData):\n",
    "    return evt.index[0], evt.index[1]\n",
    "\n",
    "\n",
    "# Implement app\n",
    "with gr.Blocks() as app:\n",
    "    gr.Markdown(\"# Interactive Remove Background from Image\")\n",
    "    with gr.Row():\n",
    "        coord_x = gr.Number(label=\"Mouse coords x\")\n",
    "        coord_y = gr.Number(label=\"Mouse coords y\")\n",
    "\n",
    "    with gr.Row():\n",
    "        input_img = gr.Image(label=\"Input image\", height=600)\n",
    "        output_img = gr.Image(label=\"Output image\", height=600)\n",
    "\n",
    "    input_img.select(get_coords, None, [coord_x, coord_y])\n",
    "    input_img.select(extract_object_by_event, [input_img], output_img)\n",
    "\n",
    "    gr.Markdown(\"## Image Examples\")\n",
    "    gr.Examples(\n",
    "        examples=[\n",
    "            [os.path.join(os.getcwd(), \"examples/dog.jpg\"), 1013, 786],\n",
    "            [os.path.join(os.getcwd(), \"examples/mannequin.jpg\"), 1720, 230],\n",
    "        ],\n",
    "        inputs=[input_img, coord_x, coord_y],\n",
    "        outputs=output_img,\n",
    "        fn=extract_object,\n",
    "        run_on_click=True,\n",
    "    )\n",
    "\n",
    "app.launch(inline=False, share=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "app.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dl_apps",
   "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.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}