create-hf-dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from Cocordiais import CocordiaisUtils as cocordiais\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",
    "DATA_INFO = [\n",
    "  {\n",
    "    \"name\": \"cordiais\",\n",
    "    \"license_id\": 1,\n",
    "    \"date\": \"2023-06-10 00:00:00\",\n",
    "    \"train_pct\": 0.5,\n",
    "    \"source_path\": os.path.join(DATA_PATH, \"cordiais-source\"),\n",
    "    \"json_path\": os.path.join(DATA_PATH, \"cordiais.json\")\n",
    "  },\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",
    "    \"name\": \"hermitage\",\n",
    "    \"license_id\": 0,\n",
    "    \"date\": \"2021-03-03 00:00:00\",\n",
    "    \"train_pct\": 1.0,\n",
    "    \"source_path\": os.path.join(DATA_PATH, \"hermitage-source\"),\n",
    "    \"json_path\": os.path.join(DATA_PATH, \"hermitage.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 = {img:id for id,img in enumerate(ALL_IMGS)}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create HF Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slug2obj = {}\n",
    "object_count = 0\n",
    "\n",
    "for split in [\"train\", \"test\"]:\n",
    "  slug2obj[split] = {}\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[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": [
    "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\"\n",
    "hf_dataset.push_to_hub(HF_DATASET, private=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test dataset"
   ]
  },
  {
   "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\""
   ]
  },
  {
   "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)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "split, idx = \"eval\", 1\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
}