create-hf-dataset-from-encord.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from Cocordiais import CocordiaisUtils as cocordiais\n",
    "\n",
    "from random import sample\n",
    "from shutil import copy2\n",
    "\n",
    "import datasets\n",
    "import json\n",
    "import os\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_PATH = \"./data\"\n",
    "\n",
    "IMG_DIRS = [\n",
    "  \"baoat-source\",\n",
    "  \"cordiais-source\",\n",
    "  \"hermitage-source\"\n",
    "]\n",
    "\n",
    "COCO_JSON_PATH = os.path.join(DATA_PATH, \"encord.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### parse images from json, add real path/source name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(COCO_JSON_PATH) as json_file_read:\n",
    "  coco_json = json.load(json_file_read)\n",
    "\n",
    "id2image = {}\n",
    "source2images = {src: [] for src in IMG_DIRS}\n",
    "\n",
    "for img in coco_json[\"images\"]:\n",
    "  img[\"file_name\"] = img[\"image_title\"]\n",
    "  del img[\"coco_url\"]\n",
    "  del img[\"image_title\"]\n",
    "\n",
    "  for ds in IMG_DIRS:\n",
    "    if os.path.isfile(os.path.join(DATA_PATH, ds, img[\"file_name\"])):\n",
    "      img[\"source\"] = ds\n",
    "      if not os.path.isfile(os.path.join(DATA_PATH, \"encord-source\", img[\"file_name\"])):\n",
    "        copy2(os.path.join(DATA_PATH, ds, img[\"file_name\"]), os.path.join(DATA_PATH, \"encord-source\"))\n",
    "      break\n",
    "\n",
    "  source2images[img[\"source\"]].append(img[\"file_name\"])\n",
    "  id2image[img[\"id\"]] = img\n",
    "\n",
    "cordiais_files = source2images[\"cordiais-source\"]\n",
    "test_size = int(0.5 * len(cordiais_files))\n",
    "test_images = cordiais_files[:test_size]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### if images already in encord-source/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(COCO_JSON_PATH) as json_file_read:\n",
    "  coco_json = json.load(json_file_read)\n",
    "\n",
    "id2image = {}\n",
    "cordiais_files = []\n",
    "\n",
    "for img in coco_json[\"images\"]:\n",
    "  img[\"file_name\"] = img[\"image_title\"]\n",
    "  del img[\"coco_url\"]\n",
    "  del img[\"image_title\"]\n",
    "\n",
    "  if os.path.isfile(os.path.join(DATA_PATH, \"cordiais-source\", img[\"file_name\"])):\n",
    "    cordiais_files.append(img[\"file_name\"])\n",
    "\n",
    "  id2image[img[\"id\"]] = img\n",
    "\n",
    "test_size = int(0.666 * len(cordiais_files))\n",
    "test_images = cordiais_files[:test_size]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slug2obj = {}\n",
    "\n",
    "for split in [\"train\", \"test\"]:\n",
    "  slug2obj[split] = {}\n",
    "\n",
    "for object_info in coco_json[\"annotations\"]:\n",
    "  img_info = id2image[object_info[\"image_id\"]]\n",
    "  file_name = img_info[\"file_name\"]\n",
    "  object_split = \"test\" if file_name in test_images else \"train\"\n",
    "\n",
    "  if file_name not in slug2obj[object_split]:\n",
    "    slug2obj[object_split][file_name] = {\n",
    "      \"image_id\": object_info[\"image_id\"],\n",
    "      \"image\": os.path.join(DATA_PATH, \"encord-source\", file_name),\n",
    "      \"image_filename\": file_name,\n",
    "      \"width\": img_info[\"width\"],\n",
    "      \"height\": img_info[\"height\"],\n",
    "      \"objects\": []\n",
    "    }\n",
    "\n",
    "  slug2obj[object_split][file_name][\"objects\"].append({\n",
    "    \"bbox_id\": object_info[\"id\"],\n",
    "    \"area\": object_info[\"area\"],\n",
    "    \"bbox\": object_info[\"bbox\"],\n",
    "    \"category\": cocordiais.ID2LABEL[object_info[\"category_id\"]],\n",
    "    \"super_category\": cocordiais.ID2SUPERLABEL[object_info[\"category_id\"]],\n",
    "    \"is_crowd\": object_info[\"iscrowd\"]\n",
    "  })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slug2obj[\"test\"][\"alberto-da-veiga-guignard_lea-e-maura.jpg\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(slug2obj[\"train\"]), len(slug2obj[\"test\"]), len(id2image), len(coco_json[\"annotations\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### add metfaces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_PATH = \"./data\"\n",
    "\n",
    "DATA_INFO = [\n",
    "  {\n",
    "    \"name\": \"metfaces\",\n",
    "    \"license_id\": 1,\n",
    "    \"date\": \"2020-06-10 00:00:00\",\n",
    "    \"train_pct\": 1.0,\n",
    "    \"source_path\": os.path.join(DATA_PATH, \"metfaces-source\"),\n",
    "    \"json_path\": os.path.join(DATA_PATH, \"metfaces.json\")\n",
    "  },\n",
    "]\n",
    "\n",
    "ALL_IMGS = []\n",
    "\n",
    "for ds in DATA_INFO:\n",
    "  file_list = sorted([f.replace(\".jpg\", \"\") for f in os.listdir(ds[\"source_path\"]) if f.endswith(\".jpg\")])\n",
    "  train_length = int(ds[\"train_pct\"] * len(file_list))\n",
    "  ds[\"source_list\"] = {}\n",
    "  ds[\"source_list\"][\"train\"] = file_list[:train_length]\n",
    "  ds[\"source_list\"][\"test\"] = file_list[train_length:]\n",
    "  ALL_IMGS += file_list\n",
    "\n",
    "img2id_met = {img:(id + len(id2image)) for id,img in enumerate(ALL_IMGS)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "object_count = len(coco_json[\"annotations\"])\n",
    "\n",
    "for ds in DATA_INFO:\n",
    "  with open(ds[\"json_path\"]) as json_file_read:\n",
    "    data_json = json.load(json_file_read)\n",
    "\n",
    "    for object_info in data_json:\n",
    "      object_slug = object_info[\"source_image\"]\n",
    "      object_split = \"train\" if object_slug in ds[\"source_list\"][\"train\"] else \"test\"\n",
    "\n",
    "      if object_slug not in slug2obj[object_split]:\n",
    "        slug2obj[object_split][object_slug] = {\n",
    "          \"image_id\": img2id_met[object_slug],\n",
    "          \"image\": os.path.join(ds[\"source_path\"], \"%s.jpg\" % object_slug),\n",
    "          \"image_filename\": \"%s.jpg\" % object_slug,\n",
    "          \"width\": object_info[\"source_image_w\"],\n",
    "          \"height\": object_info[\"source_image_h\"],\n",
    "          \"objects\": []\n",
    "        }\n",
    "\n",
    "      slug2obj[object_split][object_slug][\"objects\"].append({\n",
    "        \"bbox_id\": object_count,\n",
    "        \"area\": object_info[\"face_rect_xywh\"][2] * object_info[\"face_rect_xywh\"][3],\n",
    "        \"bbox\": object_info[\"face_rect_xywh\"],\n",
    "        \"category\": object_info[\"gender\"],\n",
    "        \"super_category\": cocordiais.LABEL2SUPERLABEL[object_info[\"gender\"]],\n",
    "        \"is_crowd\": False\n",
    "      })\n",
    "      object_count += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(slug2obj[\"train\"]), len(slug2obj[\"test\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_dict = {}\n",
    "\n",
    "for split in [\"train\", \"test\"]:\n",
    "  records = json.loads(json.dumps(list(slug2obj[split].values())))\n",
    "  for r in records:\n",
    "    r[\"objects\"] = pd.DataFrame(r[\"objects\"]).to_dict(\"list\")\n",
    "  ds_dict[split] = pd.DataFrame(records).to_dict(\"list\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hf_dataset = datasets.DatasetDict({\n",
    "  split: datasets.Dataset.from_dict(data, features=cocordiais.COCORDIAIS_FEATURES, info=cocordiais.get_dataset_info(), split=split) for split, data in ds_dict.items()\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HF_DATASET = \"thiagohersan/cordiais-faces-3\"\n",
    "hf_dataset.push_to_hub(HF_DATASET, private=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision.transforms as T\n",
    "\n",
    "from datasets import load_dataset\n",
    "from torch import ones_like\n",
    "from transformers import DetrImageProcessor\n",
    "\n",
    "from Cocordiais import CocordiaisDataset, CocordiaisUtils\n",
    "\n",
    "try:\n",
    "  HF_DATASET = HF_DATASET\n",
    "except NameError:\n",
    "  HF_DATASET = \"thiagohersan/cordiais-faces-3\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DETR_MODEL = \"facebook/detr-resnet-50\"\n",
    "\n",
    "detr_size = { \"shortest_edge\": 800, \"longest_edge\": 800 }\n",
    "detr_processor = DetrImageProcessor.from_pretrained(DETR_MODEL, size=detr_size)\n",
    "\n",
    "hf_dataset = load_dataset(HF_DATASET)\n",
    "hf_dataset_train = hf_dataset[\"train\"].train_test_split(test_size=0.2, shuffle=True, seed=101010)\n",
    "\n",
    "hf_data = {\n",
    "  \"train\": hf_dataset_train[\"train\"],\n",
    "  \"eval\": hf_dataset_train[\"test\"],\n",
    "  \"test\": hf_dataset[\"test\"]\n",
    "}\n",
    "\n",
    "dataset = {\n",
    "  \"train\": CocordiaisDataset(hf_data[\"train\"], img_processor=detr_processor, train=True),\n",
    "  \"eval\": CocordiaisDataset(hf_data[\"eval\"], img_processor=detr_processor, train=False),\n",
    "  \"test\": CocordiaisDataset(hf_data[\"test\"], img_processor=detr_processor, train=False)\n",
    "}\n",
    "\n",
    "lens = (len(dataset[\"train\"]), len(dataset[\"eval\"]), len(dataset[\"test\"]))\n",
    "\n",
    "print(\"Number of examples:\\n  Train: %s\\n  Evaluation: %s\\n  Test: %s\" % lens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "split, idx = \"train\", 0\n",
    "\n",
    "orig_image = hf_data[split][idx][\"image\"]\n",
    "detr_image = T.ToPILImage()(dataset[split].data[idx][\"pixel_values\"])\n",
    "labels = dataset[split].data[idx][\"labels\"]\n",
    "\n",
    "boxes_info = {\n",
    "  \"scores\": ones_like(labels[\"class_labels\"]),\n",
    "  \"labels\": labels[\"class_labels\"],\n",
    "  \"boxes\": CocordiaisUtils.bboxes_xcycwh_to_xyxy(labels)\n",
    "}\n",
    "\n",
    "CocordiaisUtils.plot_boxes(orig_image, boxes_info)\n",
    "CocordiaisUtils.plot_boxes(detr_image, boxes_info)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.17 ('hf-model')",
   "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.9.17"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "89e384cab7c47fb35ec95d2248b519cf922ee174880eed636c26cdfb6c4df768"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}