bg-remove.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Remove backgrounds and extract objects"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Do not remove or alter the images in the /input folder after running the object detection for this notebook to work"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import cv2\n",
    "import os.path\n",
    "from os import path as p"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Available classes in the network:\n",
      "person, bicycle, car, motorcycle, airplane, bus, train, truck, boat, traffic light, fire hydrant, stop sign, parking meter, bench, bird, cat, dog, horse, sheep, cow, elephant, bear, zebra, giraffe, backpack, umbrella, handbag, tie, suitcase, frisbee, skis, snowboard, sports ball, kite, baseball bat, baseball glove, skateboard, surfboard, tennis racket, bottle, wine glass, cup, fork, knife, spoon, bowl, banana, apple, sandwich, orange, broccoli, carrot, hot dog, pizza, donut, cake, chair, couch, potted plant, bed, dining table, toilet, tv, laptop, mouse, remote, keyboard, cell phone, microwave, oven, toaster, sink, refrigerator, book, clock, vase, scissors, teddy bear, hair drier, toothbrush\n"
     ]
    }
   ],
   "source": [
    "# input file\n",
    "FILE = \"results.pkl\"\n",
    "with open(FILE, 'rb') as fi:\n",
    "    results = pickle.load(fi)\n",
    "    \n",
    "# available classes\n",
    "classes = results['classes']\n",
    "print(\"Available classes in the network:\")\n",
    "print(*classes, sep = \", \") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting objects...\n",
      "Removed background from 'input/dogs.jpg'. Saved object in './masks/dog_0_0.png\n",
      "Removed background from 'input/dogs.jpg'. Saved object in './masks/dog_0_1.png\n",
      "Removed background from 'input/labrador.jpg'. Saved object in './masks/dog_1_0.png\n",
      "Done...\n"
     ]
    }
   ],
   "source": [
    "# Choose the object types you want to extract e.g. \"dog\"\n",
    "class_filter = [\"dog\"]\n",
    "\n",
    "if p.exists(\"masks\") is False:\n",
    "    !mkdir masks\n",
    "\n",
    "background_color = (255,255,255,0)\n",
    "\n",
    "# util method to convert the detectron2 box format\n",
    "def xyxy_to_xywh(box):\n",
    "    x1 = int(box[0])\n",
    "    y1 = int(box[1])\n",
    "    x2 = int(box[2])\n",
    "    y2 = int(box[3])\n",
    "    w = x2-x1\n",
    "    h = y2-y1\n",
    "    return [x1,y1,w,h]\n",
    "\n",
    "# Cut out masks\n",
    "print(\"Extracting objects...\")\n",
    "index = 0\n",
    "for path in results['instances']:\n",
    "    for i in range(len(results['instances'][path].pred_masks)):\n",
    "        mask_class = classes[results['instances'][path].pred_classes[i]]\n",
    "        \n",
    "        # Check if class is in filter\n",
    "        if mask_class in class_filter:\n",
    "            # make everything transparent except the mask\n",
    "            mask = results['instances'][path].pred_masks[i]\n",
    "            img = cv2.imread(path)\n",
    "            img = cv2.cvtColor(img, cv2.COLOR_RGB2RGBA)\n",
    "            x=0\n",
    "            y=0\n",
    "            for line in mask:\n",
    "                for column in line:\n",
    "                    if not column:\n",
    "                        img[x,y] = background_color\n",
    "                    y+=1\n",
    "                y=0\n",
    "                x+=1\n",
    "                \n",
    "            # Cropping image to the size of the objects bounding box\n",
    "            box = results['instances'][path].pred_boxes[i]\n",
    "            box = box.tensor.numpy()[0]\n",
    "            box = xyxy_to_xywh(box)\n",
    "            img = img[box[1]:box[1]+box[3], box[0]:box[0]+box[2]]\n",
    "            \n",
    "            # Save image\n",
    "            new_img_path = './masks/'+ mask_class+ \"_\" + str(index) + \"_\" + str(i) + \".png\"\n",
    "            cv2.imwrite(new_img_path,img)\n",
    "            print(\"Removed background from '\" + path+\"'. Saved object in '\" + new_img_path + \"\")\n",
    "    index+=1\n",
    "\n",
    "print(\"Done...\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "detectron2",
   "language": "python",
   "name": "detectron2"
  },
  "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}