notebooks/01_getting_started.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting started 🚀\n",
    "\n",
    "Welcome to the ICRA 2024 Cloth Competition! In this notebook we will load and explore the data.\n",
    "\n",
    "Run the cell below to download a part of the dataset (10 samples, this is ~1 GB) and unzip it.\n",
    "You only need to run the cell once, then you can comment it out.\n",
    "\n",
    "☁️ For the full dataset, see: https://cloud.ilabt.imec.be/index.php/s/Sy945rbamg8JMgR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Directory structure 📂"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dataclasses import fields\n",
    "from pathlib import Path\n",
    "\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import open3d as o3d\n",
    "from airo_camera_toolkit.point_clouds.conversions import point_cloud_to_open3d\n",
    "from airo_dataset_tools.data_parsers.pose import Pose\n",
    "from cloth_tools.annotation.grasp_annotation import GraspAnnotation\n",
    "from cloth_tools.dataset.format import load_competition_observation\n",
    "from cloth_tools.dataset.download import download_and_extract_dataset\n",
    "from cloth_tools.visualization.opencv import draw_pose\n",
    "\n",
    "\n",
    "data_dir = Path(\"data\")\n",
    "dataset_dir = data_dir / \"cloth_competition_dataset_0000_0-9\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.path.exists(dataset_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the cell below we download a small part (10 episodes) of the dataset if no dataset was found."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if os.path.exists(dataset_dir):\n",
    "    print(f\"Found existing dataset in: {dataset_dir}\")\n",
    "else:\n",
    "    print(f\"Downloading dataset to directory: {data_dir}\")\n",
    "    dataset_zip_url = \"https://cloud.ilabt.imec.be/index.php/s/BMg3c9g2i6oKJgN/download/cloth_competition_dataset_0000_0-9.zip\" \n",
    "    dataset_dir = download_and_extract_dataset(data_dir, dataset_zip_url)\n",
    "    dataset_dir = Path(dataset_dir)\n",
    "    print(f\"Downloaded and extracted dataset to directory: {dataset_dir}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def emoji(dir: str, file: str) -> str:\n",
    "    if os.path.isdir(os.path.join(dir, file)):\n",
    "        return \"📁\"\n",
    "    elif file.endswith(\".jpg\") or file.endswith(\".png\"):\n",
    "        return \"🖼️\"\n",
    "    elif file.endswith(\".mp4\"):\n",
    "        return \"🎥\"\n",
    "    return \"📄\"\n",
    "\n",
    "print(\"First directories in the dataset:\")\n",
    "for f in sorted(os.listdir(dataset_dir))[:5]:\n",
    "    print(emoji(dataset_dir, f), f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One sample in the dataset corresponds to one episode. \n",
    "An episode consists of one attempt at unfolding a piece of hanging cloth by grasping it at a human-annotated point.\n",
    "\n",
    "A sample directory contains the following files:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_dir = dataset_dir / \"sample_000000\"\n",
    "\n",
    "for f in os.listdir(sample_dir):\n",
    "    print(emoji(sample_dir, f), f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One sample thus contains two observations, the **start** and **result**, a **grasp** annotation and a video of the episode.\n",
    "\n",
    "* 🔎 The **start** observation is taken after the cloth has been grasped by its lowest point.\n",
    "* 🔎 The **result** observation is taken after the attempt to unfold it.\n",
    "* 👉 The **grasp** pose annotation used to unfold the garment, currently these are human-annotated.\n",
    "* 🎥 The **video** of the entire episode.\n",
    "\n",
    "Participants of the ICRA 2024 Cloth Competition will be asked to predict a good **grasp**, given the **start** observation.\n",
    "\n",
    "The grasp will be evaluated based on the **result** observation (using the surface area of cloth)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Start Observation 🔎\n",
    "\n",
    "In this section we explore some of the data contained in the start observation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "observation_start_dir = sample_dir / \"observation_start\"\n",
    "\n",
    "observation = load_competition_observation(observation_start_dir)\n",
    "\n",
    "print(\"Overview of the fields in an Cloth Competition Observation:\")\n",
    "for field in fields(observation):\n",
    "    field_name = field.name + \":\"\n",
    "    field_value = getattr(observation, field.name)\n",
    "    if isinstance(field_value, np.ndarray):\n",
    "        print(f\" - {field_name:<34} np.ndarray {field_value.shape} {field_value.dtype}\")\n",
    "    else:\n",
    "        print(f\" - {field_name:<34} {field.type}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Color images 🌈\n",
    "\n",
    "The dataset is collect using a [Zed2i](https://store.stereolabs.com/en-eu/products/zed-2i) stereo RGB-D camera 📷📷.\n",
    "For this reason, we provide two color images. \n",
    "One for the left camera and one for the right camera."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(20, 10))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.imshow(observation.image_left)\n",
    "plt.title(\"Left image\")\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.imshow(observation.image_right)\n",
    "plt.title(\"Right image\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Depth and confidence maps 🌌\n",
    "\n",
    "The slight difference in perspective between the left and right view is used by the ZED SDK to estimate depth.\n",
    "The [ZED SDK](https://www.stereolabs.com/docs/depth-sensing/depth-settings) has several depth modes, we use the NEURAL mode and enable FILL."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "depth_map = observation.depth_map\n",
    "confidence_map = observation.confidence_map\n",
    "\n",
    "print(f\"depth_map: {depth_map.shape} {depth_map.dtype}, range: {depth_map.min():.2f}-{depth_map.max():.2f}\")\n",
    "print(f\"confidence_map: {confidence_map.shape} {confidence_map.dtype}, range: {confidence_map.min():.2f}-{confidence_map.max():.2f}\")\n",
    "\n",
    "plt.figure(figsize=(20, 10))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.imshow(observation.depth_map)\n",
    "plt.title(\"Depth map\")\n",
    "plt.colorbar(fraction=0.025, pad=0.04)\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.imshow(observation.confidence_map)\n",
    "plt.title(\"Confidence map\")\n",
    "plt.colorbar(fraction=0.025, pad=0.04)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Stereo camera parameters 📷📷"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with np.printoptions(precision=3, suppress=True):\n",
    "    print(\"Resolution:\", observation.camera_resolution)\n",
    "    print(\"\\nIntrinsics (camera image formation characteristics): \\n\", observation.camera_intrinsics)\n",
    "    print(\"\\nExtrinsics (pose of the left camera in the world frame): \\n\", observation.camera_pose_in_world)\n",
    "    print(\"\\nPose of right camera expressed in left camera frame: \\n\", observation.right_camera_pose_in_left_camera)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.4 Colored point cloud ✨"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "point_cloud = observation.point_cloud\n",
    "\n",
    "point_cloud.points.shape, point_cloud.points.dtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "point_cloud.colors.shape, point_cloud.colors.dtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pcd = point_cloud_to_open3d(point_cloud)\n",
    "\n",
    "world_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.5, origin=[0, 0, 0])\n",
    "\n",
    "o3d.visualization.draw_geometries([pcd.to_legacy(), world_frame])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Grasp Annotation 👉"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grasp_dir = sample_dir / \"grasp\"\n",
    "\n",
    "for f in os.listdir(grasp_dir):\n",
    "    print(emoji(grasp_dir, f), f)\n",
    "\n",
    "\n",
    "image_frontal_annotated = cv2.imread(str(grasp_dir / \"frontal_image_grasp.jpg\"))\n",
    "image_top_annotated = cv2.imread(str(grasp_dir / \"topdown_image_grasp.jpg\"))\n",
    "\n",
    "plt.figure(figsize=(20, 10))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.imshow(cv2.cvtColor(image_frontal_annotated, cv2.COLOR_BGR2RGB))\n",
    "plt.title(\"Grasp annotation window - frontal image\")\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.imshow(cv2.cvtColor(image_top_annotated, cv2.COLOR_BGR2RGB))\n",
    "plt.title(\"Grasp annotation window - (virtual) topdown image\")\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grasp_pose_file = grasp_dir / \"grasp_pose.json\"\n",
    "grasp_annotation_file = grasp_dir / \"grasp_annotation.json\"\n",
    "\n",
    "with open(grasp_pose_file, \"r\") as f:\n",
    "    grasp_pose = Pose.model_validate_json(f.read()).as_homogeneous_matrix()\n",
    "\n",
    "\n",
    "with open(grasp_annotation_file, \"r\") as f:\n",
    "    grasp_annotation = GraspAnnotation.model_validate_json(f.read())\n",
    "\n",
    "with np.printoptions(precision=3, suppress=True):\n",
    "    print(\"Grasp pose:\\n\", grasp_pose)\n",
    "    print(\"\\nGrasp annotation:\\n\", grasp_annotation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Result Observation 🎉"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "observation_result_dir = sample_dir / \"observation_result\"\n",
    "\n",
    "observation_result = load_competition_observation(observation_result_dir)\n",
    "\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.imshow(observation_result.image_left)\n",
    "plt.title(\"Result: image of cloth after grasping and stretching\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ℹ️ The precise calculation of the evaluation metric will be released at a later date."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Coordinate frames 📐"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_W_C = observation.camera_pose_in_world\n",
    "X_W_TCPL = observation.arm_left_tcp_pose_in_world\n",
    "X_W_TCPR = observation.arm_right_tcp_pose_in_world\n",
    "X_W_LB = observation.arm_left_pose_in_world\n",
    "X_W_RB = observation.arm_right_pose_in_world\n",
    "intrinsics = observation.camera_intrinsics\n",
    "\n",
    "X_W_GRASP = grasp_pose\n",
    "\n",
    "image_bgr = cv2.cvtColor(observation.image_left, cv2.COLOR_RGB2BGR)\n",
    "\n",
    "draw_pose(image_bgr, np.identity(4), intrinsics, X_W_C, 0.25)\n",
    "draw_pose(image_bgr, X_W_LB, intrinsics, X_W_C)\n",
    "draw_pose(image_bgr, X_W_RB, intrinsics, X_W_C)\n",
    "draw_pose(image_bgr, X_W_TCPL, intrinsics, X_W_C, 0.05)\n",
    "draw_pose(image_bgr, X_W_TCPR, intrinsics, X_W_C, 0.05)\n",
    "draw_pose(image_bgr, X_W_GRASP, intrinsics, X_W_C, 0.05)\n",
    "\n",
    "image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.imshow(image_rgb)\n",
    "plt.title(\"Coordinate frames visualization\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "❔ If you have any questions, feel free to ask in on the [Github Discussions page](https://github.com/Victorlouisdg/cloth-competition/discussions)!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cloth-competition",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}