 archived/GT_semantic_segmentation_to_COCO.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Semantic Segmentation Job output conversion to COCO format\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook. \n",
    "\n",
    "![This us-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-west-2/ground_truth_labeling_jobs|ground_truth_conversion_scripts|GT_semantic_segmentation_to_COCO.ipynb)\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook walks through converting the output from SageMaker Ground Truth Semantic segmentation task into Common Objects in Context (COCO) format. The output manifest of the semantic segmentation task contains a reference to a PNG file masks for the objects that has been annotated and saved in an Amazon S3 bucket. In this notebook, I will download a sample manifest file, convert the mask into a pixelated format as in the COCO format."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install Pre-requisites"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install opencv-python pycocotools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get the output.manifest file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import os\n",
    "\n",
    "JOBNAME = \"<LABELING JOB NAME>\"  # Replace it with the labeling job name\n",
    "REGION = \"<REGION>\"  # Replace it with the job region\n",
    "client = boto3.client(\"sagemaker\", region_name=REGION)\n",
    "response = client.describe_labeling_job(LabelingJobName=JOBNAME)\n",
    "file = response[\"LabelingJobOutput\"][\"OutputDatasetS3Uri\"]\n",
    "output_manifest = os.path.basename(file)\n",
    "!aws s3 cp $file ./"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read the manifest file line by line and create \"Images\" key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "data_objs = []\n",
    "input_files = []\n",
    "images = []\n",
    "file_name = os.path.join(\"./\", output_manifest)\n",
    "\n",
    "# Assuming all images are reshaped to have the same dimensions\n",
    "height = 3956\n",
    "width = 5280\n",
    "\n",
    "# Get data objects from the output manifest\n",
    "with open(file_name) as out_manifest:\n",
    "    for line in out_manifest:\n",
    "        data_objs.append(json.loads(line))\n",
    "\n",
    "# Get the input images filenames\n",
    "for line in data_objs:\n",
    "    input_files.append(os.path.join(*line[\"source-ref\"].split(\"/\")[3:]))\n",
    "\n",
    "data_bucket = line[\"source-ref\"].split(\"/\")[2]\n",
    "dir_path = os.path.dirname(data_objs[0][\"source-ref\"])\n",
    "images_key = {\n",
    "    \"coco_url\": \"\",\n",
    "    \"date_captured\": \"\",\n",
    "    \"flickr_url\": \"\",\n",
    "    \"license\": 0,\n",
    "    \"id\": 0,\n",
    "    \"file_name\": \"\",\n",
    "    \"height\": height,\n",
    "    \"width\": width,\n",
    "}\n",
    "\n",
    "for img_id, input_file in enumerate(input_files):\n",
    "    images_key[\n",
    "        \"file_name\"\n",
    "    ] = input_file  # saving full path except bucket name. For sagemaker training, add /opt/ml/input/data prefix to the filename\n",
    "    images_key[\"id\"] = img_id\n",
    "    images.append(images_key.copy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Identify and Get Categories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "categories = []\n",
    "classes = []\n",
    "classnamesids = []\n",
    "ids = []\n",
    "names = []\n",
    "mapping = {}\n",
    "cat_hex_map = []\n",
    "\n",
    "category = {\"id\": \"\", \"name\": \"\", \"supercategory\": \"\"}\n",
    "\n",
    "for key in data_objs[0][JOBNAME + \"-ref-metadata\"][\"internal-color-map\"]:\n",
    "    classname = data_objs[0][JOBNAME + \"-ref-metadata\"][\"internal-color-map\"][key][\"class-name\"]\n",
    "    hexcolor = data_objs[0][JOBNAME + \"-ref-metadata\"][\"internal-color-map\"][key][\"hex-color\"]\n",
    "    if classname == \"BACKGROUND\":\n",
    "        continue\n",
    "    else:\n",
    "        classnamesids.append((key, classname))\n",
    "        mapping.update({hexcolor: key})\n",
    "\n",
    "for idd, classname in classnamesids:\n",
    "    category[\"id\"] = idd\n",
    "    category[\"name\"] = classname\n",
    "    classes.append(classname)\n",
    "    categories.append(category.copy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Other static variables:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "\n",
    "date_now = datetime.today().strftime(\"%Y-%m-%d\")\n",
    "year = datetime.today().strftime(\"%Y\")\n",
    "\n",
    "licenses = [{\"name\": \"\", \"id\": 0, \"url\": \"\"}]\n",
    "info = {\n",
    "    \"contributor\": \"\",\n",
    "    \"date_created\": date_now,\n",
    "    \"description\": \"ENTER YOUR DESCRIPTION HERE\",\n",
    "    \"url\": \"\",\n",
    "    \"version\": 3,\n",
    "    \"year\": \"2020\",\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Function to generate segmentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "from pycocotools import mask\n",
    "from skimage import measure\n",
    "import cv2\n",
    "from skimage import io\n",
    "\n",
    "\n",
    "def generate_segmentation(img, category_id, idd, image_id):\n",
    "    seg_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)\n",
    "    seg_img = np.uint8(seg_img)\n",
    "    f_bmask = np.asfortranarray(seg_img)\n",
    "    encoded_GT = mask.encode(f_bmask)\n",
    "    area_GT = mask.area(encoded_GT)\n",
    "    bb_GT = mask.toBbox(encoded_GT)\n",
    "    contours = measure.find_contours(seg_img, 0.5)\n",
    "    annotation = {\n",
    "        \"category_id\": category_id,\n",
    "        \"id\": idd,\n",
    "        \"image_id\": image_id,\n",
    "        \"iscrowd\": 0,\n",
    "        \"segmentation\": [],\n",
    "        \"area\": area_GT.tolist(),\n",
    "        \"bbox\": bb_GT.tolist(),\n",
    "    }\n",
    "    for contour in contours:\n",
    "        contour = np.flip(contour, axis=1)\n",
    "        segmentation = contour.ravel().tolist()\n",
    "        annotation[\"segmentation\"].append(segmentation)\n",
    "    return annotation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download the mask output from S3\n",
    "\n",
    "### WARNING: This cell will download the annotated masks from S3 bucket.\n",
    "\n",
    "If you have a very large dataset, you can either generate the segmentation data from masks using batches or make sure you have large enough storage to cater for the downloaded data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import time\n",
    "\n",
    "path = \"/\".join(data_objs[0][JOBNAME + \"-ref\"].split(\"/\")[:-1]) + \"/\"\n",
    "!aws s3 cp --recursive $path ./output_mask/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate segmentation data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import matplotlib\n",
    "from skimage import io\n",
    "import matplotlib.pyplot as plt\n",
    "from skimage.color import rgb2gray\n",
    "\n",
    "idd = 0\n",
    "annotations = []\n",
    "for image_id, file in enumerate(data_objs):\n",
    "    segments = []\n",
    "    img_file = os.path.join(\"output_mask/\", os.path.basename(file[JOBNAME + \"-ref\"]))\n",
    "    a = io.imread(img_file, plugin=\"matplotlib\")\n",
    "    colors = np.unique(a.reshape(-1, a.shape[2]), axis=0)\n",
    "    for i in range(colors.shape[0]):\n",
    "        if colors[i][:-1].mean() < 1:\n",
    "            segments.append(colors[i])\n",
    "\n",
    "    for idds, seg in enumerate(range(len(segments))):\n",
    "        print(\"Processing image {}:\".format(img_file))\n",
    "        color_hex = matplotlib.colors.to_hex(segments[seg])\n",
    "        category_id = mapping[color_hex]\n",
    "        img = io.imread(img_file, plugin=\"matplotlib\")\n",
    "        msk = segments[seg]\n",
    "        masked_img = cv2.inRange(img, msk, msk)\n",
    "        if len(segments) > 1:\n",
    "            img[masked_img > 0] = (255, 255, 255, 255)\n",
    "        annot = generate_segmentation(img, category_id, idd, image_id)\n",
    "        idd += 1\n",
    "        annotations.append(annot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Construct the input COCO file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "COCO_json = {\n",
    "    \"licenses\": licenses,\n",
    "    \"info\": info,\n",
    "    \"categories\": categories,\n",
    "    \"images\": images,\n",
    "    \"annotations\": annotations,\n",
    "}\n",
    "with open(\"COCO_file.json\", \"w\") as json_file:\n",
    "    json.dump(COCO_json, json_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test and Visualize Each Unique Category Segmentation\n",
    "\n",
    "\n",
    "This section is based on this notebook https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoDemo.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "from pycocotools.coco import COCO\n",
    "import numpy as np\n",
    "import skimage.io as io\n",
    "import matplotlib.pyplot as plt\n",
    "import pylab\n",
    "import boto3\n",
    "\n",
    "s3 = boto3.client(\"s3\")\n",
    "\n",
    "\n",
    "pylab.rcParams[\"figure.figsize\"] = (20.0, 10.0)\n",
    "annFile = \"COCO_file.json\"\n",
    "coco = COCO(annFile)\n",
    "\n",
    "# display COCO categories and supercategories\n",
    "cats = coco.loadCats(coco.getCatIds())\n",
    "nms = [cat[\"name\"] for cat in cats]\n",
    "print(\"COCO categories: \\n{}\\n\".format(\" \".join(nms)))\n",
    "nms = set([cat[\"supercategory\"] for cat in cats])\n",
    "print(\"COCO supercategories: \\n{}\".format(\" \".join(nms)))\n",
    "catIds = coco.getCatIds(catNms=classes)\n",
    "imgIds = coco.getImgIds()\n",
    "img = coco.loadImgs(imgIds[np.random.randint(0, len(imgIds))])[0]\n",
    "\n",
    "# Create local path\n",
    "folder = os.path.join(os.getcwd(), \"test_images\", *img[\"file_name\"].split(\"/\")[:-1])\n",
    "if not os.path.exists(folder):\n",
    "    os.makedirs(folder)\n",
    "\n",
    "# Download the image in the local directory\n",
    "test_image = os.path.join(\"s3://\", data_bucket, img[\"file_name\"])\n",
    "!aws s3 cp $test_image $folder\n",
    "\n",
    "# Display the image\n",
    "I = io.imread(os.path.join(\"./test_images\", img[\"file_name\"]))\n",
    "h, w, c = I.shape\n",
    "plt.axis(\"off\")\n",
    "plt.imshow(I)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Print out the annotations from the converted COCO format"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(I)\n",
    "plt.axis(\"off\")\n",
    "annIds = coco.getAnnIds(imgIds=img[\"id\"], catIds=catIds, iscrowd=None)\n",
    "anns = coco.loadAnns(annIds)\n",
    "coco.showAnns(anns)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook CI Test Results\n",
    "\n",
    "This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.\n",
    "\n",
    "![This us-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-east-1/ground_truth_labeling_jobs|ground_truth_conversion_scripts|GT_semantic_segmentation_to_COCO.ipynb)\n",
    "\n",
    "![This us-east-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-east-2/ground_truth_labeling_jobs|ground_truth_conversion_scripts|GT_semantic_segmentation_to_COCO.ipynb)\n",
    "\n",
    "![This us-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-west-1/ground_truth_labeling_jobs|ground_truth_conversion_scripts|GT_semantic_segmentation_to_COCO.ipynb)\n",
    "\n",
    "![This ca-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ca-central-1/ground_truth_labeling_jobs|ground_truth_conversion_scripts|GT_semantic_segmentation_to_COCO.ipynb)\n",
    "\n",
    "![This sa-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/sa-east-1/ground_truth_labeling_jobs|ground_truth_conversion_scripts|GT_semantic_segmentation_to_COCO.ipynb)\n",
    "\n",
    "![This eu-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-1/ground_truth_labeling_jobs|ground_truth_conversion_scripts|GT_semantic_segmentation_to_COCO.ipynb)\n",
    "\n",
    "![This eu-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-2/ground_truth_labeling_jobs|ground_truth_conversion_scripts|GT_semantic_segmentation_to_COCO.ipynb)\n",
    "\n",
    "![This eu-west-3 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-3/ground_truth_labeling_jobs|ground_truth_conversion_scripts|GT_semantic_segmentation_to_COCO.ipynb)\n",
    "\n",
    "![This eu-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-central-1/ground_truth_labeling_jobs|ground_truth_conversion_scripts|GT_semantic_segmentation_to_COCO.ipynb)\n",
    "\n",
    "![This eu-north-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-north-1/ground_truth_labeling_jobs|ground_truth_conversion_scripts|GT_semantic_segmentation_to_COCO.ipynb)\n",
    "\n",
    "![This ap-southeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-southeast-1/ground_truth_labeling_jobs|ground_truth_conversion_scripts|GT_semantic_segmentation_to_COCO.ipynb)\n",
    "\n",
    "![This ap-southeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-southeast-2/ground_truth_labeling_jobs|ground_truth_conversion_scripts|GT_semantic_segmentation_to_COCO.ipynb)\n",
    "\n",
    "![This ap-northeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-northeast-1/ground_truth_labeling_jobs|ground_truth_conversion_scripts|GT_semantic_segmentation_to_COCO.ipynb)\n",
    "\n",
    "![This ap-northeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-northeast-2/ground_truth_labeling_jobs|ground_truth_conversion_scripts|GT_semantic_segmentation_to_COCO.ipynb)\n",
    "\n",
    "![This ap-south-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-south-1/ground_truth_labeling_jobs|ground_truth_conversion_scripts|GT_semantic_segmentation_to_COCO.ipynb)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_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.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}