{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "249ca06b-e9ce-49f1-ae7b-871b43a45684",
   "metadata": {},
   "source": [
    "# AR Proof of Concept\n",
    "Inspiriert von [augmented-reality](\"https://github.com/juangallostra/augmented-reality\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be4a102b-1537-4784-ab66-c661a0136ac1",
   "metadata": {},
   "source": [
    "## Imports and Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3fb7f2c1-cb5b-4fe0-b76c-eded193d62b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Defaulting to user installation because normal site-packages is not writeable\n",
      "Requirement already satisfied: ipywidgets in /usr/lib/python3/dist-packages (6.0.0)\n",
      "Requirement already satisfied: numpy in /usr/lib/python3/dist-packages (1.21.5)\n",
      "Requirement already satisfied: opencv-python in /home/s/.local/lib/python3.10/site-packages (4.5.5.64)\n"
     ]
    }
   ],
   "source": [
    "!pip install ipywidgets numpy opencv-python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a07c9270-59df-4e61-a1d8-4752a2c912e8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "import numpy as np\n",
    "from IPython.display import display, Image\n",
    "import ipywidgets as widgets\n",
    "import threading\n",
    "from time import sleep\n",
    "import math\n",
    "import os\n",
    "\n",
    "path_list = !pwd # open in the folder where this notebook as well as models and references is located\n",
    "path = path_list[0]\n",
    "DEFAULT_COLOR = (0, 0, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "072bcec0-0c44-42ff-8414-bf36f2863dd1",
   "metadata": {},
   "source": [
    "## Helpers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b9e222b-c1f9-46fd-bcbe-fecd370b741b",
   "metadata": {},
   "source": [
    "### Button und Slider"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b954981a-7d64-48a1-8b06-c9e5dba23b08",
   "metadata": {},
   "outputs": [],
   "source": [
    "stopButton = widgets.ToggleButton(\n",
    "    value=False,\n",
    "    description='Stop',\n",
    "    disabled=False,\n",
    "    button_style='danger', # 'success', 'info', 'warning', 'danger' or ''\n",
    "    tooltip='Description',\n",
    "    icon='square' # (FontAwesome names without the `fa-` prefix)\n",
    ")\n",
    "recButton = widgets.ToggleButton(\n",
    "    value=False,\n",
    "    description='Rectangle',\n",
    "    disabled=False,\n",
    "    button_style='danger', # 'success', 'info', 'warning', 'danger' or ''\n",
    "    tooltip='Description',\n",
    "    icon='square' # (FontAwesome names without the `fa-` prefix)\n",
    ")\n",
    "matchButton = widgets.ToggleButton(\n",
    "    value=False,\n",
    "    description='Matches',\n",
    "    disabled=False,\n",
    "    button_style='danger', # 'success', 'info', 'warning', 'danger' or ''\n",
    "    tooltip='Description',\n",
    "    icon='square' # (FontAwesome names without the `fa-` prefix)\n",
    ")\n",
    "\n",
    "homog_slider = widgets.FloatSlider(value = 5, base = 10, min = 0, max = 50, step = 0.2, description = \"Homography Treshold\")\n",
    "match_slider = widgets.IntSlider(value = 0, min = 0, max = 300, description = \"Match Threshold\")\n",
    "scale_slider = widgets.IntSlider(value = 1, min = 0, max = 5, description = \"Scale\")\n",
    "layout_space = widgets.Output(layout = widgets.Layout(width = \"5px\", height = \"600px\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bb30e0b-b597-4d99-88af-61d0eac08d4c",
   "metadata": {},
   "source": [
    "### OBJ File Loader\n",
    "\n",
    "Was ist ein OBJ file? Hier ein Beispiel:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da373d52",
   "metadata": {},
   "source": [
    "```\n",
    "# List of geometric vertices, with (x, y, z [,w]) coordinates, w is optional and defaults to 1.0.\n",
    "v 0.123 0.234 0.345 1.0\n",
    "v ...\n",
    "...\n",
    "# List of texture coordinates, in (u, [,v ,w]) coordinates, these will vary between 0 and 1. v, w are optional and default to 0.\n",
    "vt 0.500 1 [0]\n",
    "vt ...\n",
    "...\n",
    "# List of vertex normals in (x,y,z) form; normals might not be unit vectors.\n",
    "vn 0.707 0.000 0.707\n",
    "vn ...\n",
    "...\n",
    "# Parameter space vertices in ( u [,v] [,w] ) form; free form geometry statement ( see below )\n",
    "vp 0.310000 3.210000 2.100000\n",
    "vp ...\n",
    "...\n",
    "# Polygonal face element (see below)\n",
    "f 1 2 3\n",
    "f 3/1 4/2 5/3\n",
    "f 6/4/1 3/5/3 7/6/5\n",
    "f 7//1 8//2 9//3\n",
    "f ...\n",
    "...\n",
    "# Line element (see below)\n",
    "l 5 8 1 2 4 9\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "820fddcc-7fbf-4131-a087-eb947874d2f0",
   "metadata": {},
   "source": [
    "https://en.wikipedia.org/wiki/Wavefront_.obj_file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5ded2809-8d44-4779-915d-cc807444fae4",
   "metadata": {},
   "outputs": [],
   "source": [
    "class OBJ:\n",
    "    def __init__(self, filename, swapyz=False):\n",
    "        \"\"\"Loads a Wavefront OBJ file. \"\"\"\n",
    "        self.vertices = []\n",
    "        self.normals = []\n",
    "        self.texcoords = []\n",
    "        self.faces = []\n",
    "        material = None\n",
    "        for line in open(filename, \"r\"):\n",
    "            if line.startswith('#'): continue\n",
    "            values = line.split()\n",
    "            if not values: continue\n",
    "            if values[0] == 'v':\n",
    "                v = list(map(float, values[1:4]))\n",
    "                if swapyz:\n",
    "                    v = v[0], v[2], v[1]\n",
    "                self.vertices.append(v)\n",
    "            elif values[0] == 'vn':\n",
    "                v = list(map(float, values[1:4]))\n",
    "                if swapyz:\n",
    "                    v = v[0], v[2], v[1]\n",
    "                self.normals.append(v)\n",
    "            elif values[0] == 'vt':\n",
    "                self.texcoords.append(map(float, values[1:3]))\n",
    "            elif values[0] == 'f':\n",
    "                face = []\n",
    "                texcoords = []\n",
    "                norms = []\n",
    "                for v in values[1:]:\n",
    "                    w = v.split('/')\n",
    "                    face.append(int(w[0]))\n",
    "                    if len(w) >= 2 and len(w[1]) > 0:\n",
    "                        texcoords.append(int(w[1]))\n",
    "                    else:\n",
    "                        texcoords.append(0)\n",
    "                    if len(w) >= 3 and len(w[2]) > 0:\n",
    "                        norms.append(int(w[2]))\n",
    "                    else:\n",
    "                        norms.append(0)\n",
    "                self.faces.append((face, norms, texcoords))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86ce6093-e35e-4abb-83e4-4283e8b931ca",
   "metadata": {},
   "source": [
    "### Render"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "16935461-fb9c-44a3-a600-0bee63d9ed47",
   "metadata": {},
   "outputs": [],
   "source": [
    "def render(img, obj, projection, model, scale, color=False):\n",
    "    \"\"\"\n",
    "    Render a loaded obj model into the current video frame\n",
    "    \"\"\"\n",
    "    vertices = obj.vertices\n",
    "    scale_matrix = np.eye(3) * scale\n",
    "    h, w = model.shape\n",
    "\n",
    "    for face in obj.faces:\n",
    "        face_vertices = face[0]\n",
    "        points = np.array([vertices[vertex - 1] for vertex in face_vertices])\n",
    "        points = np.dot(points, scale_matrix)\n",
    "        # render model in the middle of the reference surface. To do so,\n",
    "        # model points must be displaced\n",
    "        points = np.array([[p[0] + w / 2, p[1] + h / 2, p[2]] for p in points])\n",
    "        dst = cv2.perspectiveTransform(points.reshape(-1, 1, 3), projection)\n",
    "        imgpts = np.int32(dst)\n",
    "        if color is False:\n",
    "            cv2.fillConvexPoly(img, imgpts, DEFAULT_COLOR)\n",
    "        else:\n",
    "            color = hex_to_rgb(face[-1])\n",
    "            color = color[::-1]  # reverse\n",
    "            cv2.fillConvexPoly(img, imgpts, color)\n",
    "    return img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4d43240e-557e-4597-af4f-0932c3784ed0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def projection_matrix(camera_parameters, homography):\n",
    "    \"\"\"\n",
    "    From the camera calibration matrix and the estimated homography\n",
    "    compute the 3D projection matrix\n",
    "    \"\"\"\n",
    "    # Compute rotation along the x and y axis as well as the translation\n",
    "    homography = homography * (-1)\n",
    "    rot_and_transl = np.dot(np.linalg.inv(camera_parameters), homography)\n",
    "    col_1 = rot_and_transl[:, 0]\n",
    "    col_2 = rot_and_transl[:, 1]\n",
    "    col_3 = rot_and_transl[:, 2]\n",
    "    # normalise vectors\n",
    "    l = math.sqrt(np.linalg.norm(col_1, 2) * np.linalg.norm(col_2, 2))\n",
    "    rot_1 = col_1 / l\n",
    "    rot_2 = col_2 / l\n",
    "    translation = col_3 / l\n",
    "    # compute the orthonormal basis\n",
    "    c = rot_1 + rot_2\n",
    "    p = np.cross(rot_1, rot_2)\n",
    "    d = np.cross(c, p)\n",
    "    rot_1 = np.dot(c / np.linalg.norm(c, 2) + d / np.linalg.norm(d, 2), 1 / math.sqrt(2))\n",
    "    rot_2 = np.dot(c / np.linalg.norm(c, 2) - d / np.linalg.norm(d, 2), 1 / math.sqrt(2))\n",
    "    rot_3 = np.cross(rot_1, rot_2)\n",
    "    # finally, compute the 3D projection matrix from the model to the current frame\n",
    "    projection = np.stack((rot_1, rot_2, rot_3, translation)).T\n",
    "    return np.dot(camera_parameters, projection)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d1b24bd-839f-4327-a4d9-5fd7f94115f2",
   "metadata": {},
   "source": [
    "### Helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "bc42829b-a05b-43c0-9d79-d336229d5913",
   "metadata": {},
   "outputs": [],
   "source": [
    "def hex_to_rgb(hex_color):\n",
    "    \"\"\"\n",
    "    Helper function to convert hex strings to RGB\n",
    "    \"\"\"\n",
    "    hex_color = hex_color.lstrip('#')\n",
    "    h_len = len(hex_color)\n",
    "    return tuple(int(hex_color[i:i + h_len // 3], 16) for i in range(0, h_len, h_len // 3))\n",
    "\n",
    "def show_frame(frame, handle):\n",
    "    _, frame = cv2.imencode('.jpeg', frame)\n",
    "    handle.update(Image(data=frame.tobytes()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2b1a9d94-f6b6-4a8d-a06d-a443267b1cb7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def main():\n",
    "    camera_parameters = np.array([[400, 0, 320], [0, 400, 240], [0, 0, 1]])\n",
    "    # create ORB keypoint detector\n",
    "    orb = cv2.ORB_create()\n",
    "    # create BruteForceMatcher object based on hamming distance  \n",
    "    bf = cv2.BFMatcher(cv2.NORM_HAMMING, crossCheck=True)\n",
    "    # load the reference surface that will be searched in the video stream\n",
    "    dir_name = path\n",
    "    model = cv2.imread(os.path.join(dir_name, 'reference/reference.jpg'), 0)\n",
    "    # Compute model keypoints and its descriptors\n",
    "    kp_model, des_model = orb.detectAndCompute(model, None)\n",
    "    # Load 3D model from OBJ file\n",
    "    obj = OBJ(os.path.join(dir_name, 'models/fox.obj'), swapyz=True)  \n",
    "    # init video capture\n",
    "    cap = cv2.VideoCapture(0)\n",
    "    display_handle=display(None, display_id=True)\n",
    "\n",
    "    while True:\n",
    "        if stopButton.value==True:\n",
    "            cap.release()\n",
    "            display_handle.update(None)\n",
    "            return\n",
    "\n",
    "\n",
    "        sleep(0.05)\n",
    "        ret, frame = cap.read()\n",
    "        #frame = cv2.flip(frame, 1)\n",
    "        try:\n",
    "            # find and draw the keypoints of the frame\n",
    "            kp_frame, des_frame = orb.detectAndCompute(frame, None)\n",
    "            # match frame descriptors with model descriptors\n",
    "            matches = bf.match(des_model, des_frame)\n",
    "            # sort them in the order of their distance\n",
    "            # the lower the distance, the better the match\n",
    "            matches = sorted(matches, key=lambda x: x.distance)\n",
    "\n",
    "            # differenciate between source points and destination points\n",
    "            src_pts = np.float32([kp_model[m.queryIdx].pt for m in matches]).reshape(-1, 1, 2)\n",
    "            dst_pts = np.float32([kp_frame[m.trainIdx].pt for m in matches]).reshape(-1, 1, 2)\n",
    "            # compute Homography\n",
    "\n",
    "            homography, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, homog_slider.value)\n",
    "        \n",
    "            if recButton.value:\n",
    "                # Draw a rectangle that marks the found model in the frame\n",
    "                h, w = model.shape\n",
    "                pts = np.float32([[0, 0], [0, h - 1], [w - 1, h - 1], [w - 1, 0]]).reshape(-1, 1, 2)\n",
    "                # project corners into frame\n",
    "                dst = cv2.perspectiveTransform(pts, homography)\n",
    "                # connect them with lines  \n",
    "                frame = cv2.polylines(frame, [np.int32(dst)], True, 255, 3, cv2.LINE_AA)  \n",
    "            # if a valid homography matrix was found render cube on model plane\n",
    "            if match_slider.value <= len(matches):\n",
    "                if homography is not None:\n",
    "                    try:\n",
    "                        projection = projection_matrix(camera_parameters, homography)  \n",
    "                        frame = render(frame, obj, projection, model, scale_slider.value, False)\n",
    "                    except Exception:\n",
    "                        pass\n",
    "            if matchButton.value:\n",
    "                frame = cv2.drawMatches(model, kp_model, frame, kp_frame, matches, 0, flags=2)\n",
    "        except Exception:\n",
    "            pass\n",
    "        show_frame(frame, display_handle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ad6bafab",
   "metadata": {},
   "outputs": [],
   "source": [
    "thread = threading.Thread(target=main)\n",
    "thread.start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "806cb9ab-eee8-4df0-83cd-3410a63bc029",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Widget Javascript not detected.  It may not be installed or enabled properly. Reconnecting the current kernel may help.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5e3fd9c2b24e45589ddac6e8ba6f5d2b"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Widget Javascript not detected.  It may not be installed or enabled properly. Reconnecting the current kernel may help.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f53e74d8e6004936ad9e94f9da273831"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Widget Javascript not detected.  It may not be installed or enabled properly. Reconnecting the current kernel may help.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5febf1c237044eb9aec26cb0f3a0e398"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Widget Javascript not detected.  It may not be installed or enabled properly. Reconnecting the current kernel may help.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7d0bc54c6f274be3bae4238c57ffb66a"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Widget Javascript not detected.  It may not be installed or enabled properly. Reconnecting the current kernel may help.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c2d5bc425d1a4daf8df1a4ea4f8187a5"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Widget Javascript not detected.  It may not be installed or enabled properly. Reconnecting the current kernel may help.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "34cfee3383de47da8cfaefe874ead726"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Widget Javascript not detected.  It may not be installed or enabled properly. Reconnecting the current kernel may help.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "83c12d4b9e9d43e88d770fbf1650d06b"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "None"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(stopButton, homog_slider, matchButton, recButton, match_slider, scale_slider, layout_space)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8eb6803",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba5af395",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}