notebooks/hierarchy_inference.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-10T08:19:11.454412300Z",
     "start_time": "2024-07-10T08:19:11.384659200Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-10T08:19:19.147150500Z",
     "start_time": "2024-07-10T08:19:11.442118800Z"
    }
   },
   "outputs": [],
   "source": [
    "from IPython.display import display\n",
    "\n",
    "from penai.hierarchy_generation.inference import HierarchyInferencer\n",
    "from penai.hierarchy_generation.utils import (\n",
    "    InteractiveHTMLHierarchyVisualizer,\n",
    "    InteractiveSVGHierarchyVisualizer,\n",
    ")\n",
    "from penai.llm.llm_model import RegisteredLLM\n",
    "from penai.registries.projects import SavedPenpotProject\n",
    "from penai.utils.ipython import IFrameFromSrc\n",
    "from penai.utils.vis import (\n",
    "    DesignElementVisualizer,\n",
    "    ShapeHighlighter,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hierarchy Generation\n",
    "\n",
    "In this notebook we will demonstrate how to automatically infer a hierarchy of shapes with semantic shape descriptions for a Penpot project with vision language models (VLMs).\n",
    "\n",
    "First, we will load an example project and select a frame / board from a page for hierarchy generation. The current approach works on frame, respectively board level to reduce the number of shapes in a single prompt but also as boards within Penpot are typically supposed to act as logical separations of sub-designs within a single page and therefore can serve as point of reference for the LLM.\n",
    "\n",
    "Note, that the hierarchy generation works for some files and designs better than others. If a design inherently has a clear and hierarchical structure, our inference algorithm will do a pretty good job transferring this visual information into a formal structure. In cases with little inherent structure, e.g. a grid of icons, the generated hierarchies might be flat or of little information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-10T08:19:30.149660900Z",
     "start_time": "2024-07-10T08:19:19.293312400Z"
    }
   },
   "outputs": [],
   "source": [
    "project = SavedPenpotProject.MATERIAL_DESIGN_3.load(pull=True)\n",
    "cover_page = project.get_main_file().get_page_by_name(\"Cover\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we perform two important steps: removal of invisible elements and bounding box derivation. The first one is important as invisible shapes such as pure group elements that don't correspond to any visible elements can't be visually recognized by the VLM. The bounding box derivation is necessary to construct \"snippets\" of rendered elements that will be provided each separately for guiding the hierarchy generation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-10T08:19:40.815981400Z",
     "start_time": "2024-07-10T08:19:30.151903600Z"
    }
   },
   "outputs": [],
   "source": [
    "cover_page.svg.remove_elements_with_no_visible_content()\n",
    "cover_page.svg.retrieve_and_set_view_boxes_for_shape_elements()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally we will retrieve the \"Cover\" board which is the only frame in this document and covers the whole page."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-10T08:19:40.974091600Z",
     "start_time": "2024-07-10T08:19:40.817257Z"
    }
   },
   "outputs": [],
   "source": [
    "cover_frame = cover_page.svg.get_shape_by_name(\"Cover\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To now perform the hierarchy generation, we will instantiate a `HierarchyInferencer` object with a LLM of our choice and pass the prepared shape to its `infer_shape_hierarchy()`-method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-10T08:20:42.204256700Z",
     "start_time": "2024-07-10T08:19:40.957715400Z"
    }
   },
   "outputs": [],
   "source": [
    "# The ShapeHighlighter is used to create annotated visualizations of single shaped\n",
    "shape_visualizer = ShapeHighlighter()\n",
    "\n",
    "# The DesignElementVisualizer is used to visualize a design element (e.g. a primitive, group, etc.) within\n",
    "# its context in the design document.\n",
    "# It will use the ShapeHighlighter to derive visualizations for the different shapes that make up the design element.\n",
    "design_element_visualizer = DesignElementVisualizer(shape_visualizer=shape_visualizer)\n",
    "\n",
    "hierarchy_inference = HierarchyInferencer(\n",
    "    shape_visualizer=design_element_visualizer, model=RegisteredLLM.GPT4O\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the cell above finishes without errors, it indicates that the hierarchy has been derived successfully. The underlying code performs a validation of the AI response to ensure that the response format is correct (i.e. syntactically correct JSON) but also that the generated hierarchy is valid, i.e. all shapes are covered and no duplicate shapes are present.\n",
    "\n",
    "We can finally use the `InteractiveSVGHierarchyVisualizer` utility-class to visualize the generated hierarchy interactively within this notebook:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We will use the infer_shape_hierarchy_impl() method as it will return all artifacts\n",
    "# generated for the prompt, including the used visualizations and the prompt itself.\n",
    "result = hierarchy_inference.infer_shape_hierarchy_impl(cover_frame)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Optional: Display Prompt\n",
    "\n",
    "Uncomment the following lines to display the prompt that has been used to generate hierarchy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualizer = PromptVisualizer()\n",
    "# visualizer.display_message(result.prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "hierarchy = result.hierarchy\n",
    "hierarchy_svg_visualizer = InteractiveSVGHierarchyVisualizer(hierarchy, cover_frame)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-10T09:51:36.410464900Z",
     "start_time": "2024-07-10T09:51:35.724011700Z"
    }
   },
   "outputs": [],
   "source": [
    "hierarchy_html_visualizer = InteractiveHTMLHierarchyVisualizer(\n",
    "    hierarchy, svg=hierarchy_svg_visualizer.svg\n",
    ")\n",
    "display(IFrameFromSrc(hierarchy_html_visualizer.html_content, width=1200, height=900))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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": 4
}