convert the ohsome2label geococo annotation to tdml format

ohsome2label_geococo_to_tdml.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b3190c5f-86ac-488f-8979-55128e2fe839",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pytdml\n",
    "import yaml\n",
    "import json\n",
    "import geojson\n",
    "from pycocotools.coco import COCO\n",
    "from shapely.geometry import shape, box\n",
    "from pytdml.type.extended_types import EOTrainingDataset, EOTrainingData, ObjectLabel,EOTask, PixelLabel\n",
    "from pytdml.type.basic_types import Task\n",
    "from pytdml.io import write_to_json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4ff631ac-4dfa-410d-90ab-7affa4d824e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "json_file = r'data\\example_result\\annotations\\geococo.json'\n",
    "config_file = r'data\\example_config\\config.yaml'\n",
    "color_file = r'data\\example_result\\other\\colors'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "85529ea2-49a2-4607-90cb-39af49af30f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ExtendedEOTDS(EOTrainingDataset):\n",
    "    dltask: Task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2b154902-5aaa-4c44-8f4e-6ca10399dcab",
   "metadata": {},
   "outputs": [],
   "source": [
    "def COCOAnnoToInstanceLabel(coco, anno):\n",
    "    coordinates = anno[\"segmentation\"]\n",
    "    \n",
    "    label_class = coco.loadCats(anno[\"category_id\"])[0]['name']\n",
    "    \n",
    "    polygon = geojson.Polygon([coordinates])\n",
    "    \n",
    "    feature = geojson.Feature(geometry=polygon)\n",
    "    feature[\"properties\"][\"id\"] = anno[\"id\"]\n",
    "    return ObjectLabel(type =\"AI_ObjectLabel\", object=feature, label_class=label_class)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5ae0227f-3e68-41aa-97fd-f17ac3be034d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def COCOAnnoToObjectLabel(coco, anno):\n",
    "    label_class = coco.loadCats(anno[\"category_id\"])[0]['name']\n",
    "    feature = geojson.Feature(geometry = box(*anno['bbox']))\n",
    "    feature[\"properties\"][\"id\"] = anno[\"id\"]\n",
    "    return ObjectLabel(type =\"AI_ObjectLabel\", object=feature, label_class=label_class,bboxType=\"Horizental BBox\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "67ab74bc-e91d-49df-a1f9-3a145c63f654",
   "metadata": {},
   "outputs": [],
   "source": [
    "def COCOAnnoToPixelLabel(data_URL):    \n",
    "    return PixelLabel(type =\"AI_PixelLabel\", imageURL=[\"labels/\"+data_URL], imageFormat=['png'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "52b76394-7829-4fd5-8623-e151e0dbd072",
   "metadata": {},
   "outputs": [],
   "source": [
    "def COCOCatToTDMLClass(coco):\n",
    "    classes = []\n",
    "    for i in coco.getCatIds():\n",
    "        classes.append(coco.loadCats(i)[0][\"name\"])\n",
    "    return classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a7737bd3-68e7-44ad-b9ee-0705864463ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "def hex_to_rgb(hex_color):\n",
    "    hex_color = hex_color.lstrip('#')\n",
    "    r_hex = hex_color[0:2]\n",
    "    g_hex = hex_color[2:4]\n",
    "    b_hex = hex_color[4:6]\n",
    "    r = int(r_hex, 16)\n",
    "    g = int(g_hex, 16)\n",
    "    b = int(b_hex, 16)\n",
    "    return r, g, b\n",
    "\n",
    "def convert_to_rgb_dict(color_dict):\n",
    "    rgb_dicts = []\n",
    "    for key, hex_color in color_dict.items():\n",
    "        r, g, b = hex_to_rgb(hex_color)\n",
    "        rgb_string = f\"RGB({r}, {g}, {b})\"\n",
    "        rgb_dicts.append({key: rgb_string})\n",
    "    return rgb_dicts\n",
    "\n",
    "def load_color(color_file):\n",
    "    with open(color_file, 'r') as f:\n",
    "        color_dict = json.load(f)\n",
    "    return (convert_to_rgb_dict(color_dict))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "de7a1fd6-56a6-474c-a4b1-7a4ff840fc78",
   "metadata": {},
   "outputs": [],
   "source": [
    "def GeoCOCO2TDML(config_file, geococo_file, dltask,color_file=\"\"):\n",
    "    if dltask not in [\"InstanceSegmentation\", 'SemanticSegmentation',\"ObjectDetection\"]:       \n",
    "        raise NotImplementedError(rf\"{dltask} dltask is not yet implemented.\")\n",
    "        \n",
    "    with open(config_file, 'r') as f:\n",
    "        conf = yaml.safe_load(f)\n",
    "    geococo = COCO(geococo_file)\n",
    "    id = conf[\"project\"][\"name\"]\n",
    "    numOfClasses = len(geococo.getCatIds())    \n",
    "    amountOfTD = len(geococo.getImgIds())\n",
    "    eotd =[]\n",
    "    for i in geococo.getImgIds():\n",
    "        image_URL = geococo.loadImgs(i)[0][\"file_name\"]\n",
    "        if dltask == \"InstanceSegmentation\":\n",
    "            annotations = geococo.loadAnns(geococo.getAnnIds(imgIds=i))\n",
    "            tasks = Task(type=\"DL_Task\",\n",
    "                id= \"InstanceSegmentation\"\n",
    "            )            \n",
    "            classes = COCOCatToTDMLClass(geococo)\n",
    "            labels = [COCOAnnoToInstanceLabel(geococo, _) for _ in annotations]\n",
    "\n",
    "        elif dltask == \"SemanticSegmentation\":\n",
    "            tasks = Task(type=\"DL_Task\",\n",
    "                id= \"SemanticSegmentation\"\n",
    "            )\n",
    "            labels = [COCOAnnoToPixelLabel(image_URL)]\n",
    "            classes = load_color(color_file)\n",
    "\n",
    "        elif dltask == \"ObjectDetection\":\n",
    "            annotations = geococo.loadAnns(geococo.getAnnIds(imgIds=i))\n",
    "            tasks = Task(type=\"DL_Task\",\n",
    "                    id= \"ObjectDetection\"\n",
    "                )\n",
    "            classes = COCOCatToTDMLClass(geococo)\n",
    "            labels = [COCOAnnoToObjectLabel(geococo, _) for _ in annotations]\n",
    "        \n",
    "        eotd.append(EOTrainingData(type=\"AI_EOTrainingData\", \n",
    "                                   id = i,\n",
    "                                   data_URL=[\"images/\"+image_URL],\n",
    "                                   labels = labels))\n",
    "        \n",
    "    eotds = ExtendedEOTDS(type=\"AI_EOTrainingDataset\",\n",
    "                             providers=[\"ohsome2label\"],\n",
    "                             amountOfTrainingData=amountOfTD,\n",
    "                             numberOfClasses=numOfClasses,\n",
    "                             data=eotd, \n",
    "                             id=id,\n",
    "                             name=id, \n",
    "                             dltask=tasks,\n",
    "                             bands=[\"RGB\"],\n",
    "                             extent=conf['osm']['bboxes'],\n",
    "                             description=\"\",\n",
    "                             license=\"\")\n",
    "    \n",
    "    return eotds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "744798c5-56e9-45ae-8e63-7ccd3546cb1a",
   "metadata": {},
   "source": [
    "# Transfer label to TDML according to dltask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1c5f64ab-f19c-48c9-b13e-ca8d03e9728a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading annotations into memory...\n",
      "Done (t=0.01s)\n",
      "creating index...\n",
      "index created!\n",
      "loading annotations into memory...\n",
      "Done (t=0.01s)\n",
      "creating index...\n",
      "index created!\n",
      "loading annotations into memory...\n",
      "Done (t=0.00s)\n",
      "creating index...\n",
      "index created!\n"
     ]
    }
   ],
   "source": [
    "write_to_json(GeoCOCO2TDML(config_file,json_file,\"SemanticSegmentation\",color_file),\"data/example_result/pixel_label.json\")\n",
    "write_to_json(GeoCOCO2TDML(config_file,json_file,\"ObjectDetection\",color_file),\"data/example_result/object_bbox_label.json\")\n",
    "write_to_json(GeoCOCO2TDML(config_file,json_file,\"InstanceSegmentation\",color_file),\"data/example_result/instance_label.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c847b8a5-b198-4770-9bba-d09502ae57bc",
   "metadata": {},
   "outputs": [
    {
     "ename": "NotImplementedError",
     "evalue": "SomeOtherDLTask dltask is not yet implemented.",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNotImplementedError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[11], line 1\u001b[0m\n\u001b[1;32m----> 1\u001b[0m write_to_json(\u001b[43mGeoCOCO2TDML\u001b[49m\u001b[43m(\u001b[49m\u001b[43mconfig_file\u001b[49m\u001b[43m,\u001b[49m\u001b[43mjson_file\u001b[49m\u001b[43m,\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mSomeOtherDLTask\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43mcolor_file\u001b[49m\u001b[43m)\u001b[49m,\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdata/example_result/instance_label.json\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "Cell \u001b[1;32mIn[9], line 3\u001b[0m, in \u001b[0;36mGeoCOCO2TDML\u001b[1;34m(config_file, geococo_file, dltask, color_file)\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mGeoCOCO2TDML\u001b[39m(config_file, geococo_file, dltask,color_file\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n\u001b[0;32m      2\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m dltask \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m [\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mInstanceSegmentation\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mSemanticSegmentation\u001b[39m\u001b[38;5;124m'\u001b[39m,\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mObjectDetection\u001b[39m\u001b[38;5;124m\"\u001b[39m]:       \n\u001b[1;32m----> 3\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mNotImplementedError\u001b[39;00m(\u001b[38;5;124mrf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mdltask\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m dltask is not yet implemented.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m      5\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mopen\u001b[39m(config_file, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mr\u001b[39m\u001b[38;5;124m'\u001b[39m) \u001b[38;5;28;01mas\u001b[39;00m f:\n\u001b[0;32m      6\u001b[0m         conf \u001b[38;5;241m=\u001b[39m yaml\u001b[38;5;241m.\u001b[39msafe_load(f)\n",
      "\u001b[1;31mNotImplementedError\u001b[0m: SomeOtherDLTask dltask is not yet implemented."
     ]
    }
   ],
   "source": [
    "write_to_json(GeoCOCO2TDML(config_file,json_file,\"SomeOtherDLTask\",color_file),\"data/example_result/instance_label.json\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}