{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model to cpu\n"
     ]
    }
   ],
   "source": [
    "from 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 = 'cpu'\n",
    "device = 'gpu' if torch.cuda.is_available() else 'cpu'\n",
    "model_path='weights/icon_detect/best.pt'\n",
    "model_path='weights/icon_detect_v1_5/model_v1_5.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": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Florence2LanguageForConditionalGeneration has generative capabilities, as `prepare_inputs_for_generation` is explicitly overwritten. However, it doesn't directly inherit from `GenerationMixin`. From 👉v4.50👈 onwards, `PreTrainedModel` will NOT inherit from `GenerationMixin`, and this model will lose the ability to call `generate` and other related functions.\n",
      "  - If you're using `trust_remote_code=True`, you can get rid of this warning by loading the model with an auto class. See https://huggingface.co/docs/transformers/en/model_doc/auto#auto-classes\n",
      "  - If you are the owner of the model architecture code, please modify your model class such that it inherits from `GenerationMixin` (after `PreTrainedModel`, otherwise you'll get an exception).\n",
      "  - If you are not the owner of the model architecture class, please contact the model code owner to update it.\n"
     ]
    }
   ],
   "source": [
    "# two choices for caption model: fine-tuned blip2 or florence2\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",
    "# caption_model_processor = get_caption_model_processor(model_name=\"blip2\", model_name_or_path=\"weights/icon_caption_blip2\", device=device)\n",
    "caption_model_processor = get_caption_model_processor(model_name=\"florence2\", model_name_or_path=\"weights/icon_caption_florence\", device=device)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(device(type='cuda', index=0), ultralytics.models.yolo.model.YOLO)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "som_model.device, type(som_model) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "image size: (1919, 1079)\n",
      "\n",
      "image 1/1 /home/yadonglu/OmniParser/imgs/word.png: 736x1280 115 icons, 13.7ms\n",
      "Speed: 5.5ms preprocess, 13.7ms inference, 1.6ms postprocess per image at shape (1, 3, 736, 1280)\n",
      "len(filtered_boxes): 151 65\n",
      "time to prepare bbox: 0.01561737060546875\n",
      "time to process image + tokenize text inputs: 0.09026336669921875\n",
      "time to generate: 0.7382848262786865\n",
      "time to get parsed content: 0.8477945327758789\n",
      "ocr time: 0.6952385902404785\n",
      "caption time: 1.245499849319458\n"
     ]
    }
   ],
   "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",
    "# image_path = 'imgs/mobile.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.5}, 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",
    "print('ocr time:', cur_time_ocr - start)\n",
    "print('caption time:', cur_time_caption - cur_time_ocr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "image size: (1919, 1079)\n",
      "\n",
      "image 1/1 /home/yadonglu/OmniParser/imgs/word.png: 736x1280 115 icons, 299.2ms\n",
      "Speed: 5.7ms preprocess, 299.2ms inference, 3.7ms postprocess per image at shape (1, 3, 736, 1280)\n",
      "len(filtered_boxes): 151 65\n",
      "time to prepare bbox: 0.016057729721069336\n",
      "time to process image + tokenize text inputs: 1.802201509475708\n",
      "time to generate: 61.352588415145874\n",
      "time to get parsed content: 63.17377543449402\n",
      "ocr time: 0.8477699756622314\n",
      "caption time: 64.17442154884338\n"
     ]
    }
   ],
   "source": [
    "# run on cpu!!!\n",
    "# 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",
    "# image_path = 'imgs/mobile.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.5}, 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",
    "print('ocr time:', cur_time_ocr - start)\n",
    "print('caption time:', cur_time_caption - cur_time_ocr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f51b97914f0>"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 233344 characters]",
      "text/plain": [
       "<Figure size 1500x1500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "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": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>type</th>\n",
       "      <th>bbox</th>\n",
       "      <th>interactivity</th>\n",
       "      <th>content</th>\n",
       "      <th>ID</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>text</td>\n",
       "      <td>[0.023970818147063255, 0.012974976561963558, 0...</td>\n",
       "      <td>False</td>\n",
       "      <td>AutoSave</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>text</td>\n",
       "      <td>[0.0771235004067421, 0.011121409013867378, 0.1...</td>\n",
       "      <td>False</td>\n",
       "      <td>F 9~U</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>text</td>\n",
       "      <td>[0.1500781625509262, 0.011121409013867378, 0.3...</td>\n",
       "      <td>False</td>\n",
       "      <td>Document 10.docx  General*  Last Modified: Jus...</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>text</td>\n",
       "      <td>[0.4095883369445801, 0.011121409013867378, 0.4...</td>\n",
       "      <td>False</td>\n",
       "      <td>O Search</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>text</td>\n",
       "      <td>[0.010943199507892132, 0.048192769289016724, 0...</td>\n",
       "      <td>False</td>\n",
       "      <td>File</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>146</th>\n",
       "      <td>icon</td>\n",
       "      <td>[0.27768608927726746, 0.1485075205564499, 0.28...</td>\n",
       "      <td>True</td>\n",
       "      <td>Six</td>\n",
       "      <td>146</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>147</th>\n",
       "      <td>icon</td>\n",
       "      <td>[0.9438582062721252, 0.9580937027931213, 0.995...</td>\n",
       "      <td>True</td>\n",
       "      <td>battery charge indicator</td>\n",
       "      <td>147</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>148</th>\n",
       "      <td>icon</td>\n",
       "      <td>[0.31950756907463074, 0.3229200839996338, 0.33...</td>\n",
       "      <td>True</td>\n",
       "      <td>A menu or list of options.</td>\n",
       "      <td>148</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>149</th>\n",
       "      <td>icon</td>\n",
       "      <td>[0.08737719058990479, 0.148496612906456, 0.095...</td>\n",
       "      <td>True</td>\n",
       "      <td>5,5L9,5 4.5z</td>\n",
       "      <td>149</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>150</th>\n",
       "      <td>icon</td>\n",
       "      <td>[0.7414734959602356, 0.000822930654976517, 0.7...</td>\n",
       "      <td>True</td>\n",
       "      <td>Unordered List</td>\n",
       "      <td>150</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>151 rows × 5 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "     type                                               bbox  interactivity  \\\n",
       "0    text  [0.023970818147063255, 0.012974976561963558, 0...          False   \n",
       "1    text  [0.0771235004067421, 0.011121409013867378, 0.1...          False   \n",
       "2    text  [0.1500781625509262, 0.011121409013867378, 0.3...          False   \n",
       "3    text  [0.4095883369445801, 0.011121409013867378, 0.4...          False   \n",
       "4    text  [0.010943199507892132, 0.048192769289016724, 0...          False   \n",
       "..    ...                                                ...            ...   \n",
       "146  icon  [0.27768608927726746, 0.1485075205564499, 0.28...           True   \n",
       "147  icon  [0.9438582062721252, 0.9580937027931213, 0.995...           True   \n",
       "148  icon  [0.31950756907463074, 0.3229200839996338, 0.33...           True   \n",
       "149  icon  [0.08737719058990479, 0.148496612906456, 0.095...           True   \n",
       "150  icon  [0.7414734959602356, 0.000822930654976517, 0.7...           True   \n",
       "\n",
       "                                               content   ID  \n",
       "0                                             AutoSave    0  \n",
       "1                                                F 9~U    1  \n",
       "2    Document 10.docx  General*  Last Modified: Jus...    2  \n",
       "3                                             O Search    3  \n",
       "4                                                 File    4  \n",
       "..                                                 ...  ...  \n",
       "146                                                Six  146  \n",
       "147                           battery charge indicator  147  \n",
       "148                         A menu or list of options.  148  \n",
       "149                                       5,5L9,5 4.5z  149  \n",
       "150                                     Unordered List  150  \n",
       "\n",
       "[151 rows x 5 columns]"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'type': 'icon',\n",
       " 'bbox': [0.7569739818572998,\n",
       "  0.128729909658432,\n",
       "  0.7768653035163879,\n",
       "  0.14445137977600098],\n",
       " 'interactivity': True,\n",
       " 'content': 'Select option.'}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "parsed_content_list[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "omni",
   "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.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
