002_vehicles_detection.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Vehicles Segmentation \n",
    "\n",
    "In this notebook we will use the SAM model in order to enrich the the pNeuma Vision Dataset with masks and prepare a dataset for machine learning training tasks. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ff56bbd3c54e4d519edce396f3c632dc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading readme:   0%|          | 0.00/595 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "# Load the dataset in streaming mode\n",
    "dataset = load_dataset(\"katospiegel/pneuma-vision-parquet\", streaming=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_89474/3887444327.py:31: FutureWarning: Passing literal json to 'read_json' is deprecated and will be removed in a future version. To read from a literal string, wrap it in a 'StringIO' object.\n",
      "  annotation_df = pd.read_json(annotation_json, orient='columns')\n",
      "/tmp/ipykernel_89474/3887444327.py:31: FutureWarning: Passing literal json to 'read_json' is deprecated and will be removed in a future version. To read from a literal string, wrap it in a 'StringIO' object.\n",
      "  annotation_df = pd.read_json(annotation_json, orient='columns')\n"
     ]
    }
   ],
   "source": [
    "from datasets import Features, Value, Image, Dataset\n",
    "from PIL import Image as PILImage\n",
    "import io\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import torch\n",
    "from transformers import SamModel, SamProcessor\n",
    "\n",
    "\n",
    "def crop_image(img, x, y, box_size=80):\n",
    "    half_box_size = box_size // 2\n",
    "    left = max(x - half_box_size, 0)\n",
    "    upper = max(y - half_box_size, 0)\n",
    "    right = left + box_size\n",
    "    lower = upper + box_size\n",
    "    cropped_img = img.crop((left, upper, right, lower))\n",
    "    return cropped_img\n",
    "\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = SamModel.from_pretrained(\"facebook/sam-vit-huge\").to(device)\n",
    "processor = SamProcessor.from_pretrained(\"facebook/sam-vit-huge\")\n",
    "\n",
    "\n",
    "data_list = [] \n",
    "\n",
    "for data in dataset['train'].take(2):\n",
    "\n",
    "    annotation_json = data[\"Annotation_json\"]\n",
    "    annotation_df = pd.read_json(annotation_json, orient='columns')\n",
    "\n",
    "    for index, row in annotation_df.iterrows():\n",
    "        x_img = row['x_img [px]']\n",
    "        y_img = row['y_img [px]']\n",
    "\n",
    "        # Crop Image\n",
    "        box_size = 80 # Choose pair number\n",
    "        raw_image = data[\"Image\"] #PILImage.open(data[\"image\"]).convert(\"RGB\")\n",
    "\n",
    "        cropped_image = crop_image(raw_image, x_img, y_img, box_size=80)\n",
    "\n",
    "        # Image prediction\n",
    "        inputs = processor(cropped_image, return_tensors=\"pt\").to(device)\n",
    "        image_embeddings = model.get_image_embeddings(inputs[\"pixel_values\"])\n",
    "\n",
    "\n",
    "        input_points = [[[box_size/2, box_size/2]]]\n",
    "        ## Here is where we provide the input points\n",
    "        inputs = processor(cropped_image, input_points=input_points, return_tensors=\"pt\").to(device)\n",
    "        # pop the pixel_values as they are not neded\n",
    "        inputs.pop(\"pixel_values\", None)\n",
    "        inputs.update({\"image_embeddings\": image_embeddings})\n",
    "\n",
    "        with torch.no_grad():\n",
    "            outputs = model(**inputs)\n",
    "\n",
    "        masks = processor.image_processor.post_process_masks(outputs.pred_masks.cpu(), inputs[\"original_sizes\"].cpu(), inputs[\"reshaped_input_sizes\"].cpu())\n",
    "        scores = outputs.iou_scores\n",
    "\n",
    "        highest_score_mask = get_mask_with_highest_score(masks[0], scores)\n",
    "        highest_score_mask_pil = get_mask_with_highest_score_as_pil(masks[0], scores)\n",
    "\n",
    "        # Store the mask bytes in the new data dictionary\n",
    "        feature = Image()\n",
    "        new_data = {}\n",
    "        new_data['Time [s]'] = row['Time [s]']\n",
    "        new_data['id'] = row['ID']\n",
    "        new_data['Type'] = row['Type']\n",
    "        new_data['x_img [px]'] = row['x_img [px]']\n",
    "        new_data['y_img [px]'] = row['y_img [px]']\n",
    "        new_data['Angle_img [rad]'] = row['Angle_img [rad]']\n",
    "        new_data['Frame'] = data['Frame']    \n",
    "        new_data['Image'] = feature.encode_example(cropped_image)\n",
    "        new_data['Mask'] = feature.encode_example(highest_score_mask_pil)\n",
    "        data_list.append(new_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Uploading of dataset to Hugging Face"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "features = Features({\n",
    "    'Time [s]': Value(dtype='float32'),\n",
    "    'id': Value(dtype='int32'),\n",
    "    'Type': Value(dtype='string'),\n",
    "    'x_img [px]': Value(dtype='int32'),\n",
    "    'y_img [px]': Value(dtype='int32'),\n",
    "    'Angle_img [rad]': Value(dtype='float32'),\n",
    "    'Frame': Value(dtype='string'),\n",
    "    'Image': Image(decode=True),\n",
    "    'Mask': Image(decode=True)\n",
    "})\n",
    "\n",
    "data_dict = {key: [dic[key] for dic in data_list] for key in data_list[0]}\n",
    "\n",
    "# Once all data points are prepared, create the new dataset from the list\n",
    "new_dataset = Dataset.from_dict(data_dict, features=features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a434b1f8d429474c8c17853bba957c64",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Uploading the dataset shards:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b50ecbdcbe9e45808ca2e484d663e3ee",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/273 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b5f0880863144f0a88b70c2a2cec7932",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Creating parquet from Arrow format:   0%|          | 0/3 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.11/site-packages/huggingface_hub/file_download.py:991: UserWarning: Not enough free disk space to download the file. The expected file size is: 0.00 MB. The target location /root/.cache/huggingface/hub/datasets--katospiegel--ordfts-hackathon-pneuma-vehicles-segmentation/blobs only has 0.00 MB free disk space.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "af85a3ce6cc841d7ac261629afad0fd6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "README.md:   0%|          | 0.00/30.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "CommitInfo(commit_url='https://huggingface.co/datasets/katospiegel/ordfts-hackathon-pneuma-vehicles-segmentation/commit/ae21818ef53037f331ed069f6b7a99cad96cd996', commit_message='Upload dataset', commit_description='', oid='ae21818ef53037f331ed069f6b7a99cad96cd996', pr_url=None, pr_revision=None, pr_num=None)"
      ]
     },
     "execution_count": 71,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_dataset.push_to_hub(\"katospiegel/ordfts-hackathon-pneuma-vehicles-segmentation\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Previsualization of dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "<iframe\n",
       "  src=\"https://huggingface.co/datasets/katospiegel/ordfts-hackathon-pneuma-vehicles-segmentation/embed/viewer/default/train\"\n",
       "  frameborder=\"0\"\n",
       "  width=\"100%\"\n",
       "  height=\"560px\"\n",
       "></iframe>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%html\n",
    "\n",
    "<iframe\n",
    "  src=\"https://huggingface.co/datasets/katospiegel/ordfts-hackathon-pneuma-vehicles-segmentation/embed/viewer/default/train\"\n",
    "  frameborder=\"0\"\n",
    "  width=\"100%\"\n",
    "  height=\"560px\"\n",
    "></iframe>"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}