grounded_sam_instructedit_demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/wangq0e/miniconda3/envs/eda_1/lib/python3.8/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import argparse\n",
    "import os\n",
    "import copy\n",
    "import gradio as gr\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "from PIL import Image, ImageDraw, ImageFont\n",
    "\n",
    "\n",
    "# Grounding DINO\n",
    "import GroundingDINO.groundingdino.datasets.transforms as T\n",
    "from GroundingDINO.groundingdino.models import build_model\n",
    "from GroundingDINO.groundingdino.util import box_ops\n",
    "from GroundingDINO.groundingdino.util.slconfig import SLConfig\n",
    "from GroundingDINO.groundingdino.util.utils import clean_state_dict, get_phrases_from_posmap\n",
    "\n",
    "# segment anything\n",
    "from segment_anything import build_sam, SamPredictor \n",
    "import cv2\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# diffusers\n",
    "import PIL\n",
    "import requests\n",
    "import torch\n",
    "from io import BytesIO\n",
    "from stable_diffusion_masked_diffedit import StableDiffusionMaskedDiffeditPipeline\n",
    "from diffusers import DDIMScheduler\n",
    "\n",
    "# chatgpt\n",
    "from chatgpt import call_chatgpt\n",
    "\n",
    "# blip2\n",
    "from transformers import AutoProcessor, Blip2ForConditionalGeneration\n",
    "\n",
    "from operator import itemgetter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_model(model_config_path, model_checkpoint_path, device):\n",
    "    args = SLConfig.fromfile(model_config_path)\n",
    "    args.device = device\n",
    "    model = build_model(args)\n",
    "    checkpoint = torch.load(model_checkpoint_path, map_location=\"cpu\")\n",
    "    load_res = model.load_state_dict(clean_state_dict(checkpoint[\"model\"]), strict=False)\n",
    "    print(load_res)\n",
    "    _ = model.eval()\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/wangq0e/miniconda3/envs/eda_1/lib/python3.8/site-packages/torch/functional.py:478: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at  /opt/conda/conda-bld/pytorch_1659484810403/work/aten/src/ATen/native/TensorShape.cpp:2894.)\n",
      "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "final text_encoder_type: bert-base-uncased\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertModel: ['cls.seq_relationship.weight', 'cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.LayerNorm.weight', 'cls.seq_relationship.bias', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.bias']\n",
      "- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_IncompatibleKeys(missing_keys=[], unexpected_keys=['label_enc.weight'])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading checkpoint shards: 100%|██████████| 2/2 [00:11<00:00,  5.60s/it]\n"
     ]
    }
   ],
   "source": [
    "# cfg\n",
    "config_file = 'GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py'  # change the path of the model config file\n",
    "grounded_checkpoint = 'GroundingDINO/weights/groundingdino_swint_ogc.pth'  # change the path of the model\n",
    "sam_checkpoint = 'sam_vit_h_4b8939.pth'\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "\n",
    "# load model\n",
    "model = load_model(config_file, grounded_checkpoint, device=device)\n",
    "blip_processor = AutoProcessor.from_pretrained('Salesforce/blip2-flan-t5-xl')\n",
    "blip_model = Blip2ForConditionalGeneration.from_pretrained('Salesforce/blip2-flan-t5-xl', torch_dtype=torch.float16).to(device) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_image(image_path, height=512, width=512,):\n",
    "    # load image\n",
    "    image_pil = Image.open(image_path).convert(\"RGB\")  # load image\n",
    "    image_pil = preprocess_image(image_pil, height, width)\n",
    "\n",
    "    transform = T.Compose(\n",
    "        [\n",
    "            T.RandomResize([800], max_size=1333),\n",
    "            T.ToTensor(),\n",
    "            T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),\n",
    "        ]\n",
    "    )\n",
    "    image, _ = transform(image_pil, None)  # 3, h, w\n",
    "    return image_pil, image\n",
    "\n",
    "\n",
    "def get_grounding_output(model, image, caption, box_threshold, text_threshold, with_logits=True, device=\"cpu\", return_max=False):\n",
    "    caption = caption.lower()\n",
    "    caption = caption.strip()\n",
    "    if not caption.endswith(\".\"):\n",
    "        caption = caption + \".\"\n",
    "    model = model.to(device)\n",
    "    image = image.to(device)\n",
    "    with torch.no_grad():\n",
    "        outputs = model(image[None], captions=[caption])\n",
    "    logits = outputs[\"pred_logits\"].cpu().sigmoid()[0]  # (nq, 256)\n",
    "    boxes = outputs[\"pred_boxes\"].cpu()[0]  # (nq, 4)\n",
    "    logits.shape[0]\n",
    "\n",
    "    # filter output\n",
    "    logits_filt = logits.clone()\n",
    "    boxes_filt = boxes.clone()\n",
    "    filt_mask = logits_filt.max(dim=1)[0] > box_threshold\n",
    "    logits_filt = logits_filt[filt_mask]  # num_filt, 256\n",
    "    boxes_filt = boxes_filt[filt_mask]  # num_filt, 4\n",
    "    logits_filt.shape[0]\n",
    "\n",
    "    # get phrase\n",
    "    tokenlizer = model.tokenizer\n",
    "    tokenized = tokenlizer(caption)\n",
    "    # build pred\n",
    "    pred_phrases = []\n",
    "    \n",
    "    logits_filt_item = []\n",
    "    \n",
    "    for logit, box in zip(logits_filt, boxes_filt):\n",
    "        pred_phrase = get_phrases_from_posmap(logit > text_threshold, tokenized, tokenlizer)\n",
    "        logits_filt_item.append(logit.max().item())\n",
    "        if with_logits:\n",
    "            pred_phrases.append(pred_phrase + f\"({str(logit.max().item())[:4]})\")\n",
    "        else:\n",
    "            pred_phrases.append(pred_phrase)\n",
    "    \n",
    "    if return_max:\n",
    "        max_logit_index, max_logit = max(enumerate(logits_filt_item), key=itemgetter(1))\n",
    "        return boxes_filt[max_logit_index].unsqueeze(0), [pred_phrases[max_logit_index]]\n",
    "\n",
    "    return boxes_filt, pred_phrases\n",
    "\n",
    "def show_mask(mask, ax, random_color=False, opacity=0.6):\n",
    "    if random_color:\n",
    "        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)\n",
    "    else:\n",
    "        color = np.array([30/255, 144/255, 255/255, opacity])\n",
    "    h, w = mask.shape[-2:]\n",
    "    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)\n",
    "    ax.imshow(mask_image)\n",
    "\n",
    "\n",
    "def show_box(box, ax, label):\n",
    "    x0, y0 = box[0], box[1]\n",
    "    w, h = box[2] - box[0], box[3] - box[1]\n",
    "    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2)) \n",
    "    ax.text(x0, y0, label)\n",
    "    \n",
    "def adjust_encoding_ratio(instructions, feedbacks, cur_encoding_ratio):\n",
    "    full_instructions = []\n",
    "    for instruction in instructions:\n",
    "        full_instruction = '{}. The current encoding ratio is {}.'.format(instruction, cur_encoding_ratio)\n",
    "        full_instructions.append(full_instruction)\n",
    "        \n",
    "    encoding_ratio = call_chatgpt(full_instructions, feedbacks=feedbacks)\n",
    "    \n",
    "    return encoding_ratio\n",
    "\n",
    "def preprocess_image(image, height=512, width=512, left=0, right=0, top=0, bottom=0):\n",
    "    if isinstance(image, str):\n",
    "        image = np.array(Image.open(image))\n",
    "    elif isinstance(image, np.ndarray):\n",
    "        pass\n",
    "    else:\n",
    "        image = np.array(image)\n",
    "        \n",
    "    if image.ndim == 3:\n",
    "        image = image[:, :, :3]\n",
    "        h, w, _ = image.shape\n",
    "    else:\n",
    "        h, w = image.shape\n",
    "        \n",
    "    left = min(left, w-1)\n",
    "    right = min(right, w - left - 1)\n",
    "    top = min(top, h - left - 1)\n",
    "    bottom = min(bottom, h - top - 1)\n",
    "    image = image[top:h-bottom, left:w-right]\n",
    "    \n",
    "    if image.ndim == 3:\n",
    "        h, w, _ = image.shape\n",
    "    else:\n",
    "        h, w = image.shape\n",
    "    if h < w:\n",
    "        offset = (w - h) // 2\n",
    "        image = image[:, offset:offset + h]\n",
    "    elif w < h:\n",
    "        offset = (h - w) // 2\n",
    "        image = image[offset:offset + w]\n",
    "    image = Image.fromarray(image).resize((height, width))\n",
    "    return image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def maskedit(image_path, user_instructions, is_blip2_description, encoding_ratio, mask_mode, \n",
    "             return_max=True, box_threshold=0.3, text_threshold=0.25, is_show_box=True):\n",
    "    # load image\n",
    "    image_pil, image = load_image(image_path, height=512, width=512,)    \n",
    "\n",
    "    if is_blip2_description:\n",
    "        \n",
    "        blip_prompt_1 = \"Is this a photo, a painting, a drawing, or other kind of arts?\"\n",
    "        blip_inputs = blip_processor(image_pil, text=blip_prompt_1, return_tensors=\"pt\").to(device, torch.float16)\n",
    "        generated_ids = blip_model.generate(**blip_inputs, max_new_tokens=20)\n",
    "        generated_text = blip_processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()\n",
    "        blip_prompt_2 = \"{} of\".format(generated_text.capitalize())\n",
    "        blip_inputs = blip_processor(image_pil, text=blip_prompt_2, return_tensors=\"pt\").to(device, torch.float16)\n",
    "        generated_ids = blip_model.generate(**blip_inputs, max_new_tokens=20)\n",
    "        description = blip_processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()\n",
    "        description = \"{} {}.\".format(blip_prompt_2, description)\n",
    "        descriptions = [description]    \n",
    "        print('Blip2 description: {}'.format(description))\n",
    "    else:\n",
    "        descriptions = None\n",
    "    instructions = [user_instructions]\n",
    "    det_prompt, prompt_inversion, prompt, _ = call_chatgpt(instructions, descriptions=descriptions)\n",
    "    print('Segmentation prompt: {}'.format(det_prompt))\n",
    "    print('Inverted prompt: {}'.format(prompt_inversion))\n",
    "    print('Editing prompt: {}'.format(prompt))\n",
    "    \n",
    "    prompts = [prompt_inversion, prompt]\n",
    "\n",
    "    # load model\n",
    "    model = load_model(config_file, grounded_checkpoint, device=device)\n",
    "\n",
    "    # run grounding dino model\n",
    "    boxes_filt, pred_phrases = get_grounding_output(\n",
    "        model, image, det_prompt, box_threshold, text_threshold, device=device, return_max=return_max,\n",
    "    )\n",
    "\n",
    "    # initialize SAM\n",
    "    predictor = SamPredictor(build_sam(checkpoint=sam_checkpoint).to(device))\n",
    "    image = cv2.imread(image_path)\n",
    "    image = np.asarray(preprocess_image(image, height=512, width=512))\n",
    "    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "    predictor.set_image(image)\n",
    "\n",
    "    size = image_pil.size\n",
    "    H, W = size[1], size[0]\n",
    "    for i in range(boxes_filt.size(0)):\n",
    "        boxes_filt[i] = boxes_filt[i] * torch.Tensor([W, H, W, H])\n",
    "        boxes_filt[i][:2] -= boxes_filt[i][2:] / 2\n",
    "        boxes_filt[i][2:] += boxes_filt[i][:2]\n",
    "\n",
    "    boxes_filt = boxes_filt.cpu()\n",
    "    transformed_boxes = predictor.transform.apply_boxes_torch(boxes_filt, image.shape[:2]).to(device)\n",
    "\n",
    "    masks, _, _ = predictor.predict_torch(\n",
    "        point_coords = None,\n",
    "        point_labels = None,\n",
    "        boxes = transformed_boxes.to(device),\n",
    "        multimask_output = False,\n",
    "    )\n",
    "\n",
    "    if is_show_mask:\n",
    "        # draw output image\n",
    "        plt.figure(figsize=(10, 10))\n",
    "        plt.imshow(image)\n",
    "        for mask in masks:\n",
    "            show_mask(mask.cpu().numpy(), plt.gca(), random_color=True)\n",
    "        if is_show_box:\n",
    "            for box, label in zip(boxes_filt, pred_phrases):\n",
    "                show_box(box.numpy(), plt.gca(), label)\n",
    "\n",
    "        plt.axis('off')\n",
    "        \n",
    "    # Masked diffedit pipeline\n",
    "    if mask_mode == 'merge':\n",
    "        masks = torch.sum(masks, dim=0).unsqueeze(0)\n",
    "        masks = torch.where(masks > 0, True, False)\n",
    "    mask = masks[0][0].cpu().numpy() # simply choose the first mask, which will be refine in the future release\n",
    "    mask_pil = Image.fromarray(mask)\n",
    "    image_pil = Image.fromarray(image)\n",
    "\n",
    "    scheduler =DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule=\"scaled_linear\", clip_sample=False,\n",
    "                                        set_alpha_to_one=False)\n",
    "    pipe = StableDiffusionMaskedDiffeditPipeline.from_pretrained(\"CompVis/stable-diffusion-v1-4\", scheduler=scheduler)\n",
    "    pipe.to(device)\n",
    "\n",
    "    image = pipe(prompt, image=image_pil, mask_image=mask_pil, prompt_inversion=prompt_inversion, encoding_ratio=encoding_ratio,\n",
    "                 height=512, width=512, ).images[0]\n",
    "    image = image.resize((512, 512))\n",
    "\n",
    "    return image, mask_pil, prompts, image_pil"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_name = \"cat_dog\"\n",
    "image_path = './example_input/{}.jpg'.format(image_name)\n",
    "user_instructions = \"Change the cat to a dog\"\n",
    "is_blip2_description = True\n",
    "is_show_mask = True\n",
    "encoding_ratio = 0.5\n",
    "box_threshold = 0.3\n",
    "text_threshold = 0.25\n",
    "output_image_path = \"./output_image/{}.jpg\".format(image_name)\n",
    "output_mask_path = './output_mask/{}_mask.jpg'.format(image_name)\n",
    "mask_mode = \"max\"\n",
    "if mask_mode == \"max\":\n",
    "    return_max = True\n",
    "else:\n",
    "    return_max = False\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "output_image, output_mask, prompts, input_image = maskedit(image_path, user_instructions, \n",
    "                                                           is_blip2_description, encoding_ratio, \n",
    "                                                           mask_mode, return_max=return_max, box_threshold=box_threshold, text_threshold=text_threshold)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "blobgan",
   "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.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "5b48a3a89d390d3619abb9164aeaae43be92ed2bd36c8c7dce0401e5db318047"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}