data-labelling/GT_semantic_segmentation_to_COCO.ipynb (471 lines of code) (raw):
{
"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",
"\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",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\n",
"\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
}