{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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 = 'cuda'\n",
    "\n",
    "som_model = get_yolo_model(model_path='weights/icon_detect/best.pt')\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",
    "\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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "som_model.device, type(som_model) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "cnt = 0\n",
    "image_path = 'imgs/google_page.png'\n",
    "# image_path = 'imgs/windows_home.png'\n",
    "image_path = 'imgs/windows_multitab.png'\n",
    "draw_bbox_config = {\n",
    "    'text_scale': 0.8,\n",
    "    'text_thickness': 2,\n",
    "    'text_padding': 3,\n",
    "    'thickness': 3,\n",
    "}\n",
    "BOX_TRESHOLD = 0.03\n",
    "\n",
    "image = Image.open(image_path)\n",
    "image_rgb = image.convert('RGB')\n",
    "\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})\n",
    "text, ocr_bbox = ocr_bbox_rslt\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=False, 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.1)\n",
    "\n",
    "\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=(12,12))\n",
    "\n",
    "image = Image.open(io.BytesIO(base64.b64decode(dino_labled_img)))\n",
    "plt.axis('off')\n",
    "\n",
    "plt.imshow(image)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(label_coordinates)\n",
    "parsed_content_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pilot",
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}