{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from util.utils import get_som_labeled_img, check_ocr_box, get_caption_model_processor, get_yolo_model\n",
    "import torch\n",
    "from ultralytics import YOLO\n",
    "from PIL import Image\n",
    "device = 'cuda'\n",
    "model_path='weights/OmniParser-v2/icon_detect/model.pt'\n",
    "\n",
    "som_model = get_yolo_model(model_path)\n",
    "\n",
    "som_model.to(device)\n",
    "print('model to {}'.format(device))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# two choices for caption model: fine-tuned blip2 or florence2\n",
    "import importlib\n",
    "# import util.utils\n",
    "# importlib.reload(utils)\n",
    "from util.utils import get_som_labeled_img, check_ocr_box, get_caption_model_processor, get_yolo_model\n",
    "caption_model_processor = get_caption_model_processor(model_name=\"florence2\", model_name_or_path=\"weights/OmniParser-v2/icon_caption\", device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "som_model.device, type(som_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reload utils\n",
    "import importlib\n",
    "import utils\n",
    "importlib.reload(utils)\n",
    "# from utils import get_som_labeled_img, check_ocr_box, get_caption_model_processor, get_yolo_model\n",
    "\n",
    "image_path = 'imgs/google_page.png'\n",
    "image_path = 'imgs/windows_home.png'\n",
    "# image_path = 'imgs/windows_multitab.png'\n",
    "# image_path = 'imgs/omni3.jpg'\n",
    "# image_path = 'imgs/ios.png'\n",
    "image_path = 'imgs/word.png'\n",
    "# image_path = 'imgs/excel2.png'\n",
    "\n",
    "image = Image.open(image_path)\n",
    "image_rgb = image.convert('RGB')\n",
    "print('image size:', image.size)\n",
    "\n",
    "box_overlay_ratio = max(image.size) / 3200\n",
    "draw_bbox_config = {\n",
    "    'text_scale': 0.8 * box_overlay_ratio,\n",
    "    'text_thickness': max(int(2 * box_overlay_ratio), 1),\n",
    "    'text_padding': max(int(3 * box_overlay_ratio), 1),\n",
    "    'thickness': max(int(3 * box_overlay_ratio), 1),\n",
    "}\n",
    "BOX_TRESHOLD = 0.05\n",
    "\n",
    "import time\n",
    "start = time.time()\n",
    "ocr_bbox_rslt, is_goal_filtered = check_ocr_box(image_path, display_img = False, output_bb_format='xyxy', goal_filtering=None, easyocr_args={'paragraph': False, 'text_threshold':0.9}, use_paddleocr=True)\n",
    "text, ocr_bbox = ocr_bbox_rslt\n",
    "cur_time_ocr = time.time() \n",
    "\n",
    "dino_labled_img, label_coordinates, parsed_content_list = get_som_labeled_img(image_path, som_model, BOX_TRESHOLD = BOX_TRESHOLD, output_coord_in_ratio=True, ocr_bbox=ocr_bbox,draw_bbox_config=draw_bbox_config, caption_model_processor=caption_model_processor, ocr_text=text,use_local_semantics=True, iou_threshold=0.7, scale_img=False, batch_size=128)\n",
    "cur_time_caption = time.time() \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot dino_labled_img it is in base64\n",
    "import base64\n",
    "import matplotlib.pyplot as plt\n",
    "import io\n",
    "plt.figure(figsize=(15,15))\n",
    "\n",
    "image = Image.open(io.BytesIO(base64.b64decode(dino_labled_img)))\n",
    "plt.axis('off')\n",
    "\n",
    "plt.imshow(image)\n",
    "# print(len(parsed_content_list))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "df = pd.DataFrame(parsed_content_list)\n",
    "df['ID'] = range(len(df))\n",
    "\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parsed_content_list"
   ]
  }
 ],
 "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}