promopt_mask.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#import library\n",
    "import numpy as np\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "from utils.sample_points import uniform_sampling\n",
    "# import sys\n",
    "# sys.path.append(\"..\")\n",
    "from segment_anything import sam_model_registry, SamPredictor\n",
    "from utils.sample_points import uniform_sampling\n",
    "from cityscapesscripts.helpers.labels import trainId2label as trainid2label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_mask(mask, ax, random_color=False):\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, 0.6])\n",
    "    h, w = mask.shape[-2:]\n",
    "    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)\n",
    "    ax.imshow(mask_image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_points(coords, labels, ax, marker_size=375):\n",
    "    pos_points = coords[labels==1]\n",
    "    neg_points = coords[labels==0]\n",
    "    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)\n",
    "    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)   \n",
    "\n",
    "def show_points2(coords, labels, ax, marker_size=25):\n",
    "    pos_points = coords[labels==1]\n",
    "    neg_points = coords[labels==0]\n",
    "    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)\n",
    "    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_box(box, ax):\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from segment_anything.utils.transforms import ResizeLongestSide\n",
    "resize_transform = ResizeLongestSide(sam.image_encoder.img_size)\n",
    "\n",
    "def prepare_image(image, transform, device):\n",
    "    image = transform.apply_image(image)\n",
    "    image = torch.as_tensor(image, device=device.device) \n",
    "    return image.permute(2, 0, 1).contiguous()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def trainid2color(trainid):\n",
    "    '''\n",
    "    function: convert trainID to color in cityscapes\n",
    "    input: trainid\n",
    "    output: color\n",
    "    '''\n",
    "    #if the input is a number in np.uint8, it means it is a trainid\n",
    "    if type(trainid) == np.uint8:\n",
    "        label_object = trainid2label[trainid]\n",
    "        return label_object.color[::-1]\n",
    "    else:\n",
    "        color_mask = np.zeros((trainid.shape[0], 3), dtype=np.uint8)\n",
    "        for i in range(trainid.shape[0]):\n",
    "            label_object = trainid2label[trainid[i]]\n",
    "            color_mask[i] = label_object.color[::-1]\n",
    "    return color_mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def color_segmentation(segmentation):\n",
    "    #get the color segmentation result, initial the color segmentation result with black (0,0,0)\n",
    "    #input: segmentation [h, w]\n",
    "    color_segmentation = np.zeros((segmentation.shape[0], segmentation.shape[1], 3), dtype=np.uint8)\n",
    "    train_ids = np.unique(segmentation)\n",
    "    for train_id in train_ids:\n",
    "        color_segmentation[segmentation == train_id] = trainid2color(train_id)[::-1]\n",
    "    return color_segmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thing_ids = range(19)\n",
    "class_names = ['road', 'sidewalk', 'building', 'wall', 'fence', 'pole', 'traffic light', 'traffic sign', 'vegetation', 'terrain', 'sky', 'person', 'rider', 'car', 'truck', 'bus', 'train', 'motorcycle', 'bicycle']\n",
    "# thing_ids = [11, 12, 13, 14, 15, 16, 17, 18]\n",
    "\n",
    "#load image\n",
    "image = cv2.imread('images/aachen_000029_000019_leftImg8bit.png')\n",
    "image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "label = cv2.imread('images/aachen_000029_000019_gtFine_labelTrainIds.png', cv2.IMREAD_GRAYSCALE)\n",
    "color_label = color_segmentation(label)\n",
    "\n",
    "#show image\n",
    "plt.figure(figsize=(12,6))\n",
    "plt.imshow(image)\n",
    "plt.axis('on')\n",
    "plt.show()\n",
    "plt.close()\n",
    "\n",
    "#show label\n",
    "plt.figure(figsize=(12,6))  # constrained_layout=True\n",
    "plt.imshow(color_label)\n",
    "plt.axis('off')\n",
    "# plt.tight_layout()\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load model\n",
    "sam_checkpoint = \"models/sam_vit_h_4b8939.pth\" #sam_vit_b_01ec64.pth, sam_vit_l_0b3195.pth\n",
    "model_type = \"vit_h\" #vit_b, vit_h, vit_l\n",
    "\n",
    "device = \"cuda:0\"\n",
    "\n",
    "sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)\n",
    "sam.to(device=device)\n",
    "print('load model successfully')\n",
    "\n",
    "#set predictor\n",
    "predictor = SamPredictor(sam)\n",
    "#Process the image to produce an image embedding by calling `SamPredictor.set_image`. \n",
    "#`SamPredictor` remembers this embedding and will use it for subsequent mask prediction.\n",
    "predictor.set_image(image) #get the image embedding, for subsequent mask prediction\n",
    "\n",
    "#show the points on the image\n",
    "# input_point = np.array([[500, 375]]) #input the point prompt in the format of (x, y)\n",
    "# input_label = np.array([1]) #input the label prompt in the format of (0: background, 1: foreground)\n",
    "\n",
    "# plt.figure(figsize=(10,10))\n",
    "# plt.imshow(image)\n",
    "# show_points(input_point, input_label, plt.gca())\n",
    "# plt.axis('on')\n",
    "# plt.show()  \n",
    "\n",
    "# Predict with `SamPredictor.predict`. \n",
    "# The model returns masks, quality predictions for those masks, \n",
    "# and low resolution mask logits that can be passed to the next iteration of prediction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# point-level prompt\n",
    "# input_point = np.array([[1024, 490], [1075, 482], [1001, 452]]) #车1\n",
    "# input_point = np.array([[526, 471], [529, 427], [640, 488], [690, 471]]) #车2\n",
    "# input_point = np.array([[1024, 490], [1075, 482], [1001, 452], [526, 471], [529, 427], [640, 488], [690, 471]]) #车3\n",
    "# input_point = np.array([[363, 103], [392, 541], [385, 593]]) # vegetation1\n",
    "# input_point = np.array([[988, 270], [918, 302], [1030, 324], [883, 342], [1134, 305], [1115, 253]]) #vegetation 2\n",
    "# input_point = np.array([[363, 103], [392, 541], [385, 593], [988, 270], [918, 302], [1030, 324], [883, 342], [1134, 305], [1115, 253]])\n",
    "# input_point = np.array([[668, 364], [670, 459], [668, 568]]) # pole\n",
    "# input_point = np.array([[1651, 369], [1649, 465], [1652, 568]])\n",
    "input_point = np.array([[668, 364], [670, 459], [668, 568], [1651, 369], [1649, 465], [1652, 568]])\n",
    "\n",
    "input_label = np.array([1] * len(input_point))\n",
    "# input_point = np.array([[1024, 490], [1075, 482], [1001, 452], [549, 463], [584, 432], [603, 501]])\n",
    "# input_label = np.array([1, 1, 1, 1, 1, 1])\n",
    "\n",
    "masks, scores, logits = predictor.predict(point_coords=input_point, point_labels=input_label, multimask_output=True)\n",
    "\n",
    "for i, (mask, score) in enumerate(zip(masks, scores)):\n",
    "        plt.figure(figsize=(16, 8))\n",
    "        plt.imshow(image)\n",
    "        show_mask(mask, plt.gca())\n",
    "        show_points(input_point, input_label, plt.gca())\n",
    "        plt.title(f\"Mask {i+1}, Score: {score:.3f}\", fontsize=18)\n",
    "        plt.axis('off')\n",
    "        plt.show()\n",
    "        plt.close()\n",
    "\n",
    "# input_point2 = np.array([[1024, 490], [1075, 482], [1001, 452], [549, 463], [584, 432], [603, 501]])\n",
    "# input_label2 = np.array([1, 1, 1, 1, 1, 1])\n",
    "\n",
    "masks2, scores2, logits2 = predictor.predict(point_coords=input_point, point_labels=input_label, multimask_output=False)\n",
    "plt.figure(figsize=(16, 8))\n",
    "plt.imshow(image)\n",
    "show_mask(masks2, plt.gca())\n",
    "show_points(input_point, input_label, plt.gca())\n",
    "plt.title(f\"Mask, Score: {scores2[0]:.3f}\", fontsize=18)\n",
    "plt.axis('off')\n",
    "plt.show()\n",
    "plt.close()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# box-level prompt\n",
    "image1_boxes = torch.tensor([\n",
    "    [952, 428, 1117, 553],\n",
    "    [479, 412, 720, 535],\n",
    "    [727, 428, 838, 509],\n",
    "    [1374, 400, 1416, 528],\n",
    "    [1344, 321, 1376, 364],\n",
    "    # [1361, 301, 1375, 527],\n",
    "    [1818, 464, 1987, 568],\n",
    "    [1201, 447, 1221, 488],\n",
    "], device=sam.device)\n",
    "\n",
    "batched_input = [\n",
    "     {\n",
    "         'image': prepare_image(image, resize_transform, sam),\n",
    "         'boxes': resize_transform.apply_boxes_torch(image1_boxes, image.shape[:2]),\n",
    "         'original_size': image.shape[:2]  # \n",
    "     },\n",
    "]\n",
    "\n",
    "batched_output = sam(batched_input, multimask_output=False)\n",
    "\n",
    "fig, ax = plt.subplots(1, 1, figsize=(10, 10))\n",
    "\n",
    "ax.imshow(image)\n",
    "for mask in batched_output[0]['masks']:\n",
    "    show_mask(mask.cpu().numpy(), ax, random_color=True)\n",
    "for box in image1_boxes:\n",
    "    show_box(box.cpu().numpy(), ax)\n",
    "ax.axis('off')\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "unique_ids = np.unique(label)\n",
    "unique_ids = unique_ids[unique_ids != 255]\n",
    "for unique_id in unique_ids:\n",
    "    if unique_id not in thing_ids:\n",
    "        continue\n",
    "    print('class:', class_names[unique_id])\n",
    "    \n",
    "    #foreground points\n",
    "    id_mask = label == unique_id\n",
    "    sample_points = uniform_sampling(id_mask, 0.001, 10)\n",
    "    input_point = np.array(sample_points)\n",
    "    input_label = np.array([1] * len(sample_points))\n",
    "    \n",
    "    # #add background points\n",
    "    # bg_mask = id_mask == False\n",
    "    # sample_bg_points = uniform_sampling(bg_mask, 0.0001, 10)\n",
    "    # background_point = np.array(sample_bg_points)\n",
    "    # background_label = np.array([0] * len(sample_bg_points))\n",
    "    \n",
    "    # #concate the foreground points and background points\n",
    "    # input_point = np.concatenate([input_point, background_point], axis=0)\n",
    "    # input_label = np.concatenate([input_label, background_label], axis=0)\n",
    "    \n",
    "\n",
    "    # masks, scores, logits = predictor.predict(\n",
    "    # point_coords=input_point,\n",
    "    # point_labels=input_label,\n",
    "    # multimask_output=True,\n",
    "    # )\n",
    "\n",
    "    # for i, (mask, score) in enumerate(zip(masks, scores)):\n",
    "    #     plt.figure(figsize=(16, 8))\n",
    "    #     plt.imshow(image)\n",
    "    #     show_mask(mask, plt.gca())\n",
    "    #     show_points2(input_point, input_label, plt.gca())\n",
    "    #     plt.title(f\"Mask {i+1}, Score: {score:.3f}\", fontsize=18)\n",
    "    #     plt.axis('off')\n",
    "    #     plt.savefig('outputs_tmp/{}_{}.png'.format(class_names[unique_id], i), bbox_inches='tight', pad_inches=0)\n",
    "    #     plt.show()\n",
    "    # plt.close()\n",
    "    \n",
    "# input_point = np.array([[500, 375], [1125, 625]])\n",
    "# input_label = np.array([1, 1])\n",
    "\n",
    "    # mask_input = logits[np.argmax(scores), :, :]  # Choose the model's best mask\n",
    "    \n",
    "    # masks, _, _ = predictor.predict(\n",
    "    # point_coords=input_point,\n",
    "    # point_labels=input_label,\n",
    "    # mask_input=mask_input[None, :, :],\n",
    "    # multimask_output=False,\n",
    "    # )\n",
    "    \n",
    "    masks, _, _ = predictor.predict(\n",
    "    point_coords=input_point,\n",
    "    point_labels=input_label,\n",
    "    multimask_output=False,\n",
    "    )\n",
    "\n",
    "    plt.figure(figsize=(16, 8))\n",
    "    plt.imshow(image)\n",
    "    show_mask(masks, plt.gca())\n",
    "    show_points2(input_point, input_label, plt.gca())\n",
    "    plt.title(f\"Class: {class_names[unique_id]}, Score: {score:.3f}\", fontsize=18)\n",
    "    plt.axis('off')\n",
    "    plt.savefig('outputs_tmp/{}.png'.format(class_names[unique_id]), bbox_inches='tight', pad_inches=0)\n",
    "    plt.show()\n",
    "\n",
    "# input_point = np.array([[500, 375], [1125, 625]])\n",
    "# input_label = np.array([1, 0])\n",
    "\n",
    "# mask_input = logits[np.argmax(scores), :, :]  # Choose the model's best mask\n",
    "\n",
    "# masks, _, _ = predictor.predict(\n",
    "# point_coords=input_point,\n",
    "# point_labels=input_label,\n",
    "# mask_input=mask_input[None, :, :],\n",
    "# multimask_output=False,\n",
    "# )\n",
    "\n",
    "# plt.figure(figsize=(10, 10))\n",
    "# plt.imshow(image)\n",
    "# show_mask(masks, plt.gca())\n",
    "# show_points(input_point, input_label, plt.gca())\n",
    "# plt.axis('off')\n",
    "# plt.show() "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.6.9 64-bit",
   "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.8.5"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}