scripts/notebooks/visualize_infer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HAS_CD_TO_ROOT = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "\n",
    "if HAS_CD_TO_ROOT is False:\n",
    "    os.chdir(\"../../\")\n",
    "    HAS_CD_TO_ROOT = True\n",
    "\n",
    "import logging\n",
    "import os\n",
    "from typing import Optional, Dict\n",
    "\n",
    "import hydra\n",
    "import torch\n",
    "from hydra.utils import instantiate\n",
    "from datasets import DatasetDict, load_dataset, IterableDatasetDict\n",
    "from omegaconf import DictConfig, OmegaConf\n",
    "from src.data.transforms import SamCaptionerDataTransform\n",
    "from src.data.collator import SamCaptionerDataCollator\n",
    "from src.arguments import Arguments, global_setup, SAMCaptionerModelArguments, SCAModelArguments, SCAModelBaseArguments\n",
    "from src.models.sam_captioner import SAMCaptionerConfig, SAMCaptionerModel, SAMCaptionerProcessor\n",
    "from src.models.sca import ScaProcessor\n",
    "\n",
    "from transformers.trainer_utils import get_last_checkpoint\n",
    "from transformers import set_seed, Trainer\n",
    "import gradio as gr\n",
    "from dataclasses import dataclass\n",
    "import numpy as np\n",
    "from functools import partial\n",
    "import pandas as pd\n",
    "from src.train import prepare_datasets, prepare_data_transform, prepare_processor\n",
    "import pycocotools.mask\n",
    "from PIL import Image\n",
    "\n",
    "from hydra import initialize, compose\n",
    "import json\n",
    "import tqdm\n",
    "import hashlib\n",
    "import glob\n",
    "import cv2  \n",
    "import numpy as np  \n",
    "from PIL import Image, ImageDraw, ImageFont  \n",
    "import random\n",
    "import pycocotools.mask\n",
    "import sqlite3\n",
    "from contextlib import closing\n",
    "import dotenv\n",
    "\n",
    "os.getcwd()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET='vg-densecap-local'\n",
    "with initialize(version_base=\"1.3\", config_path=\"../../src/conf\"):\n",
    "    args = compose(\n",
    "        config_name=\"conf\",\n",
    "        overrides=[\n",
    "            f\"train_data=[{DATASET}]\",\n",
    "            f\"eval_data=[{DATASET}]\",\n",
    "            \"+model=base_sam_captioner\",\n",
    "            \"training.output_dir=tmp/visualization\"\n",
    "            # \"training.do_train=True\",\n",
    "            # \"training.do_eval=True\",\n",
    "        ],\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "args, training_args, model_args = global_setup(args)\n",
    "os.makedirs(training_args.output_dir, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize our dataset and prepare it\n",
    "with initialize(version_base=\"1.3\", config_path=\"../../src/conf\"):\n",
    "    train_dataset, eval_dataset = prepare_datasets(args)\n",
    "\n",
    "# NOTE(xiaoke): load sas_key from .env for huggingface model downloading.\n",
    "dotenv.load_dotenv('.env')\n",
    "use_auth_token = os.getenv(\"USE_AUTH_TOKEN\", False)\n",
    "\n",
    "processor = prepare_processor(model_args, use_auth_token)\n",
    "\n",
    "train_dataset, eval_dataset = prepare_data_transform(\n",
    "    training_args, model_args, train_dataset, eval_dataset, processor\n",
    ")\n",
    "\n",
    "\n",
    "# [NOTE] Used to restore the image tensor after transformed\n",
    "# Use global to avoid passing too many arguments\n",
    "global image_mean, image_std\n",
    "image_mean, image_std = (\n",
    "    processor.sam_processor.image_processor.image_mean,\n",
    "    processor.sam_processor.image_processor.image_std,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "REWRITE_MAPPING = False\n",
    "image_id_to_dataset_id_mapping_file = os.path.join(training_args.output_dir, \"image_id_to_dataset_id_mapping.json\")\n",
    "\n",
    "def find_json_file_with_md5(json_file):\n",
    "    json_file_name, json_file_ext = os.path.splitext(json_file)\n",
    "    json_file_blob = f\"{json_file_name}-*{json_file_ext}\"\n",
    "    return glob.glob(json_file_blob)\n",
    "\n",
    "def get_md5_from_json(json_file):\n",
    "    with open(json_file, \"r\") as f:\n",
    "        content = f.read()\n",
    "    return hashlib.md5(content.encode()).hexdigest()\n",
    "\n",
    "def get_md5_from_pyobj(pyobj):\n",
    "    bytes_data = pyobj.encode()\n",
    "    readable_hash = hashlib.md5(bytes_data).hexdigest()  \n",
    "    return readable_hash\n",
    "\n",
    "def save_dict_to_json_with_md5(json_file, dict_data):\n",
    "    # Convert to json and bytes  \n",
    "    json_data = json.dumps(dict_data)  \n",
    "    json_data_md5 = get_md5_from_pyobj(json_data)\n",
    "    json_file_name, json_file_ext = os.path.splitext(json_file)\n",
    "    json_file_with_md5 = f\"{json_file_name}-{json_data_md5}{json_file_ext}\"\n",
    "    with open(json_file_with_md5, 'w') as f:  \n",
    "        f.write(json_data)  \n",
    "    return json_file_with_md5\n",
    "\n",
    "# Initialize our dataset and prepare it\n",
    "with initialize(version_base=\"1.3\", config_path=\"../../src/conf\"):\n",
    "    args_no_image = compose(\n",
    "        config_name=\"conf\",\n",
    "        overrides=[\n",
    "            f\"train_data=[{DATASET}]\",\n",
    "            f\"eval_data=[{DATASET}]\",\n",
    "            \"+model=base_sam_captioner\",\n",
    "            \"training.output_dir=tmp/visualization\"\n",
    "            # \"training.do_train=True\",\n",
    "            # \"training.do_eval=True\",\n",
    "        ],\n",
    "    )\n",
    "    args_no_image.train_data_overrides = ['data.with_image=False']\n",
    "    args_no_image.eval_data_overrides = ['data.with_image=False']\n",
    "    train_dataset_no_image, eval_dataset_no_image = prepare_datasets(args_no_image)\n",
    "\n",
    "json_file_with_md5_ls = find_json_file_with_md5(image_id_to_dataset_id_mapping_file)\n",
    "if len(json_file_with_md5_ls) > 1:\n",
    "    raise ValueError(f\"find more than one json file with md5, {json_file_with_md5_ls}\")\n",
    "if REWRITE_MAPPING is False and len(json_file_with_md5_ls) == 1:\n",
    "    image_id_to_dataset_id_mapping_file = json_file_with_md5_ls[0]\n",
    "    md5_in_name = os.path.splitext(image_id_to_dataset_id_mapping_file)[0].split(\"-\")[-1]\n",
    "    assert md5_in_name == get_md5_from_json(image_id_to_dataset_id_mapping_file), f\"md5 not match for {image_id_to_dataset_id_mapping_file}\"\n",
    "\n",
    "    with open(image_id_to_dataset_id_mapping_file, \"r\") as f:\n",
    "        image_id_to_dataset_id_mapping = json.load(f)\n",
    "    print(f\"Load mapping from {image_id_to_dataset_id_mapping_file}\")\n",
    "else:\n",
    "    image_id_to_dataset_id_mapping = {\n",
    "        \"train\": dict(),\n",
    "        **{k: dict() for k in eval_dataset_no_image.keys()},\n",
    "    }\n",
    "    for sample_cnt, sample in enumerate(tqdm.tqdm(train_dataset_no_image)):\n",
    "        image_id_to_dataset_id_mapping[\"train\"][sample[\"image_id\"]] = sample_cnt\n",
    "    for eval_dataset_name, eval_dataset_  in eval_dataset_no_image.items():\n",
    "        for sample_cnt, sample in enumerate(tqdm.tqdm(eval_dataset_)):\n",
    "            image_id_to_dataset_id_mapping[eval_dataset_name][sample[\"image_id\"]] = sample_cnt\n",
    "    image_id_to_dataset_id_mapping_file = save_dict_to_json_with_md5(image_id_to_dataset_id_mapping_file, image_id_to_dataset_id_mapping)\n",
    "    print(f\"save mapping to {image_id_to_dataset_id_mapping_file}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the infer json\n",
    "infer_json_path_dict = {\n",
    "    \"vg-gpt2l-bs_32-lsj\": \"/home/t-yutonglin/xiaoke/segment-caption-anything-v2/amlt/111423.exp-only_vg-finetune_vg/111323.infer-train-sca-ablat-lsj-scale_lr-110423.4x8_fin-16x4_unfin.pre/best-gpt2-large-lsj-1xlr.110423.octo-4x8-v100-16g-no_pre/vg-densecap-region_descriptions/infer-post_processed/infer-visual_genome-region_descriptions_v1.2.0-test.json\",\n",
    "    \"vg-ollm3bv2-bs_32-lsj\": \"amlt/110723.exp.ablat-lsj-scale_lr-running-2/infer-train-sca-ablat-lsj-scale_lr-110423-110723.running-2/best-fp16-ollm3bv2-large-lsj-1xlr.110423.octo-4x8-v100-16g-no_pre/vg-densecap-region_descriptions/infer-post_processed/infer-visual_genome-region_descriptions_v1.2.0-test.json\",\n",
    "    \"o365_vg-gpt2l-bs_64-lsj\": \"amlt/111423.exp-only_vg-finetune_vg/111323.infer-train-sca.finetune_lsj_scale_lr-o365_1e_4_1xlr_lsj.111023.4x8_fin-16x4_unfin.pre/best-111223.rr1-4x8-v100-32g-pre.fintune-gpt2_large-lr_1e_4-1xlr-lsj-bs_2-o365_1e_4_no_lsj_bs_64/vg-densecap-region_descriptions/infer-post_processed/infer-visual_genome-region_descriptions_v1.2.0-test.json\",\n",
    "}\n",
    "\n",
    "for job_name, json_path in infer_json_path_dict.items():\n",
    "    print(f\"job_name: {job_name}\")\n",
    "    print(f\"is exists: {os.path.exists(json_path)}\")\n",
    "    assert os.path.exists(json_path), f\"{json_path} not exists\"\n",
    "\n",
    "infer_json_path = infer_json_path_dict[\"vg-gpt2l-bs_32-lsj\"]\n",
    "with open(infer_json_path, \"r\") as f:\n",
    "    infer_json = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import colorsys  \n",
    "colors = [  \n",
    "    (235, 206, 135),  # Soft Yellow  \n",
    "    (176, 224, 230),  # Powder Blue  \n",
    "    (240, 230, 140),  # Khaki  \n",
    "    (244, 164, 96),   # Sandy Brown  \n",
    "    (144, 238, 144),  # Light Green  \n",
    "    (221, 160, 221),  # Plum  \n",
    "    (255, 182, 193),  # Light Pink  \n",
    "    (173, 216, 230),  # Light Blue  \n",
    "    (255, 235, 205),  # Blanched Almond  \n",
    "    (245, 255, 250),  # Mint Cream  \n",
    "]  \n",
    "  \n",
    "# Convert RGB to HSV and keep track of original index  \n",
    "colors_hsv = [(colorsys.rgb_to_hsv(color[0]/255, color[1]/255, color[2]/255), index) for index, color in enumerate(colors)]  \n",
    "  \n",
    "# Sort by hue  \n",
    "colors_hsv.sort()  \n",
    "  \n",
    "# Convert back to RGB  \n",
    "harmonious_colors = [colors[index] for hsv, index in colors_hsv]  \n",
    "\n",
    "# Your selected colors  \n",
    "selected_colors = harmonious_colors\n",
    "# Calculate height of each color strip  \n",
    "height = 256 // len(selected_colors)  \n",
    "  \n",
    "# Create a new image with RGB mode  \n",
    "img = Image.new('RGB', (256, 256))  \n",
    "  \n",
    "draw = ImageDraw.Draw(img)  \n",
    "  \n",
    "for i, color in enumerate(selected_colors):  \n",
    "    # Calculate the start and end positions of the color strip  \n",
    "    start_pos = i * height  \n",
    "    end_pos = start_pos + height  \n",
    "  \n",
    "    # Draw the color strip  \n",
    "    draw.rectangle([(0, start_pos), (256, end_pos)], fill=color)  \n",
    "img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hex_to_rgb(hex_color):  \n",
    "    return tuple([int(hex_color[i:i+2], 16) for i in (1, 3, 5)])\n",
    "  \n",
    "hex_colors = [\"#B0F2BCFF\", \"#89E8ACFF\", \"#67DBA5FF\", \"#4CC8A3FF\", \"#38B2A3FF\", \"#2C98A0FF\", \"#257D98FF\"]  \n",
    "  \n",
    "rgb_colors = [hex_to_rgb(color[:-2]) for color in hex_colors]  # '[:-2]' is to remove the 'FF' at the end of each color code, which represents the alpha channel in ARGB format  \n",
    "harmonious_colors = rgb_colors\n",
    "  \n",
    "# Create a new image with RGB mode  \n",
    "img = Image.new('RGB', (256, 256))  \n",
    "  \n",
    "draw = ImageDraw.Draw(img)  \n",
    "\n",
    "print(rgb_colors)  \n",
    "for i, color in enumerate(rgb_colors):  \n",
    "    # Calculate the start and end positions of the color strip  \n",
    "    start_pos = i * height  \n",
    "    end_pos = start_pos + height  \n",
    "  \n",
    "    # Draw the color strip  \n",
    "    print(color)\n",
    "    draw.rectangle([(0, start_pos), (256, end_pos)], fill=color)  \n",
    "img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "EVAL_DATASET_SPLIT = 'visual_genome-densecap-local-densecap-test'\n",
    "first_sample = infer_json[3]\n",
    "references = first_sample[\"references\"]\n",
    "candidates = first_sample[\"candidates\"]\n",
    "\n",
    "image_id = first_sample[\"metadata\"][\"metadata_image_id\"]\n",
    "region_id = first_sample[\"metadata\"][\"metadata_region_id\"]\n",
    "input_boxes = first_sample[\"metadata\"][\"metadata_input_boxes\"]\n",
    "\n",
    "sample_cnt = image_id_to_dataset_id_mapping[EVAL_DATASET_SPLIT][str(image_id)]\n",
    "sample = eval_dataset[EVAL_DATASET_SPLIT][sample_cnt]\n",
    "image = sample[\"image\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image, ImageDraw, ImageFont  \n",
    "import cv2  \n",
    "import numpy as np  \n",
    "\n",
    "FONT_PATH = \"tmp/Arial.ttf\"\n",
    "FONT = ImageFont.truetype(FONT_PATH, 20)\n",
    "\n",
    "def draw_bbox(pil_image, bbox, color=(30, 144, 255), thickness=1):  \n",
    "    cv_image = np.array(pil_image)\n",
    "    cv2.rectangle(cv_image, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, thickness)  \n",
    "    return Image.fromarray(cv_image)  \n",
    "\n",
    "def draw_mask(pil_image, mask_array, color=(30, 144, 255), alpha=0.5):  \n",
    "    cv_image = np.array(pil_image)\n",
    "    cv_image[mask_array == 1] = cv_image[mask_array == 1] * (1 - alpha) + np.array(color) * alpha\n",
    "    return Image.fromarray(cv_image)\n",
    "\n",
    "def draw_mask_boundary(pil_image, mask_array, color=(30, 144, 255), thickness=1):  \n",
    "    cv_image = np.array(pil_image)\n",
    "    contours, _ = cv2.findContours(mask_array, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n",
    "    cv2.drawContours(cv_image, contours, -1, color, thickness)\n",
    "    return Image.fromarray(cv_image)\n",
    "\n",
    "def resize_image(image, height=None, width=None):\n",
    "    \"\"\"\n",
    "    Resizes an image given the desired height and/or width.\n",
    "    If only one of height or width is provided, the other dimension is scaled proportionally.\n",
    "    If both height and width are provided, the image is resized to the exact dimensions.\n",
    "    \"\"\"\n",
    "    if height is None and width is None:\n",
    "        return image\n",
    "    \n",
    "    original_width, original_height = image.size\n",
    "    \n",
    "    if height is not None and width is not None:\n",
    "        new_size = (width, height)\n",
    "    elif height is not None:\n",
    "        new_size = (int(original_width * height / original_height), height)\n",
    "    else:\n",
    "        new_size = (width, int(original_height * width / original_width))\n",
    "    \n",
    "    return image.resize(new_size)\n",
    "\n",
    "\n",
    "\n",
    "def draw_captions(pil_image, captions, font_path, font_size=20, font_color=(0, 0, 0), bg_color=(255, 255, 255), margin_size=10, captions_color=None):  \n",
    "    font = ImageFont.truetype(font_path, font_size)  \n",
    "    # Calculate the total height of the padding for the captions  \n",
    "    total_height = 0  \n",
    "    for caption in captions:  \n",
    "        _, _, text_width, text_height = font.getbbox(caption)\n",
    "        total_height += text_height + margin_size  \n",
    "  \n",
    "    # Create a new image with padding at the bottom for the captions  \n",
    "    new_image = Image.new('RGB', (pil_image.width, pil_image.height + total_height), bg_color)  \n",
    "    new_image.paste(pil_image, (0, 0))  \n",
    "\n",
    "    draw = ImageDraw.Draw(new_image)  \n",
    "    # Draw each caption  \n",
    "    y_position = pil_image.height  \n",
    "    for caption_id, caption in enumerate(captions):  \n",
    "        _, _, text_width, text_height = font.getbbox(caption)\n",
    "        if captions_color is not None:\n",
    "            text_bbox = (0, y_position, text_width, y_position + text_height)\n",
    "            fill_color = captions_color[caption_id]\n",
    "            draw.rectangle(text_bbox, fill=fill_color, width=0)\n",
    "        draw.text((0, y_position), caption, fill=font_color, font=font)  \n",
    "        y_position += text_height + margin_size\n",
    "  \n",
    "    return new_image  \n",
    "  \n",
    "def plot_bbox_and_captions(pil_image, bbox=None, captions=None, mask=None, font_path='tmp/Arial.ttf', font_size=20, font_color=(0, 0, 0), bg_color=(255, 255, 255), margin_size=0, captions_color=None):  \n",
    "    if bbox is not None:\n",
    "        pil_image = draw_bbox(pil_image, bbox)  \n",
    "    if mask is not None:\n",
    "        pil_image = draw_mask_boundary(pil_image, mask)\n",
    "    pil_image = resize_image(pil_image, height=512)\n",
    "    if captions is not None:\n",
    "        pil_image = draw_captions(pil_image, captions, font_path, font_size, font_color, bg_color, margin_size, captions_color=captions_color)  \n",
    "    return pil_image  \n",
    "\n",
    "\n",
    "font_path = 'tmp/Arial.ttf'\n",
    "captions = candidates + references\n",
    "\n",
    "import random\n",
    "# Calculate the number of colors  \n",
    "num_colors = len(harmonious_colors)  \n",
    "# Generate a random start index  \n",
    "start_index = random.randint(0, num_colors - 1)  \n",
    "# Select colors in a round-robin way  \n",
    "selected_colors = [harmonious_colors[(start_index + i) % num_colors] for i in range(len(captions))]  \n",
    "captions_color = selected_colors\n",
    "\n",
    "pil_img_with_bbox_and_captions = plot_bbox_and_captions(image, bbox=input_boxes, captions=captions, captions_color=captions_color)  \n",
    "pil_img_with_bbox_and_captions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mask_db_file = 'tmp/sam_mask_db/visual_genome-densecap-local-densecap-test/results.db'\n",
    "with closing(sqlite3.connect(mask_db_file)) as conn:\n",
    "    cursor = conn.cursor()\n",
    "    cursor.execute(\n",
    "        \"\"\"  \n",
    "        SELECT region_cnt, image_cnt, region_id, image_id, masks, scores, input_box, gt_caption\n",
    "        FROM results where region_cnt = ?\n",
    "    \"\"\", (3,)\n",
    "    )\n",
    "    results = cursor.fetchall()\n",
    "    print(results)\n",
    "rle_masks = results[0][4]\n",
    "scores = results[0][5]\n",
    "rle_masks = json.loads(rle_masks)\n",
    "scores = json.loads(scores)\n",
    "masks = pycocotools.mask.decode(rle_masks)\n",
    "\n",
    "pil_img_with_bbox_and_captions = plot_bbox_and_captions(image, bbox=input_boxes, mask=masks[..., -1], captions=captions, captions_color=captions_color)  \n",
    "pil_img_with_bbox_and_captions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the infer json\n",
    "infer_json_path_dict = {\n",
    "    \"sam_cap-git_large\": \"amlt/111523.exp.sam_captioner/infer_sam_captioner_region_chunkify/microsoft/git-large/infer-post_processed/infer-visual_genome-densecap-local-densecap-test.json.post.json\",\n",
    "    \"sam_cap-blip_large\": \"amlt/111523.exp.sam_captioner/infer-sam_captioner-region_chunkify-eval_suite/Salesforce/blip-image-captioning-large/vg-densecap-region_descriptions/infer-post_processed/infer-visual_genome-region_descriptions_v1.2.0-test.json.post.json\",\n",
    "    \"sam_cap-blip2_opt_2_7b\": \"amlt/111523.exp.sam_captioner/infer-sam_captioner-region_chunkify-eval_suite/Salesforce/blip2-opt-2.7b/infer-post_processed/infer-visual_genome-densecap-local-densecap-test.json.post.json\",\n",
    "    \"grit\": \"amlt/111523.exp.grit/infer-promptable-grit/infer-post_processed/infer-visual_genome-densecap-local-densecap-test.json.post.json\", \n",
    "    \"vg-gpt2l-bs_32-lsj\": \"/home/t-yutonglin/xiaoke/segment-caption-anything-v2/amlt/111423.exp-only_vg-finetune_vg/111323.infer-train-sca-ablat-lsj-scale_lr-110423.4x8_fin-16x4_unfin.pre/best-gpt2-large-lsj-1xlr.110423.octo-4x8-v100-16g-no_pre/vg-densecap-region_descriptions/infer-post_processed/infer-visual_genome-region_descriptions_v1.2.0-test.json\",\n",
    "    \"vg-ollm3bv2-bs_32-lsj\": \"/home/t-yutonglin/xiaoke/segment-caption-anything-v2/amlt/110723.exp.ablat-lsj-scale_lr-running-2/infer-train-sca-ablat-lsj-scale_lr-110423-110723.running-2/best-fp16-ollm3bv2-large-lsj-1xlr.110423.octo-4x8-v100-16g-no_pre/vg-densecap-region_descriptions/infer-post_processed/infer-visual_genome-region_descriptions_v1.2.0-test.json\",\n",
    "    \"o365_vg-gpt2l-bs_64-lsj\": \"/home/t-yutonglin/xiaoke/segment-caption-anything-v2/amlt/111423.exp-only_vg-finetune_vg/111323.infer-train-sca.finetune_lsj_scale_lr-o365_1e_4_1xlr_lsj.111023.4x8_fin-16x4_unfin.pre/best-111223.rr1-4x8-v100-32g-pre.fintune-gpt2_large-lr_1e_4-1xlr-lsj-bs_2-o365_1e_4_no_lsj_bs_64/vg-densecap-region_descriptions/infer-post_processed/infer-visual_genome-region_descriptions_v1.2.0-test.json\",\n",
    "}\n",
    "\n",
    "for job_name, json_path in infer_json_path_dict.items():\n",
    "    print(f\"job_name: {job_name}\")\n",
    "    print(f\"is exists: {os.path.exists(json_path)}\")\n",
    "    assert os.path.exists(json_path), f\"{json_path} not exists\"\n",
    "\n",
    "class MultiInferJson(torch.utils.data.Dataset):\n",
    "    def __init__(self, infer_json_path_dict):\n",
    "        self.infer_json_path_dict = infer_json_path_dict\n",
    "        self.infer_json_dict = dict()\n",
    "        for job_name, json_path in tqdm.tqdm(self.infer_json_path_dict.items(), desc=\"Load json\"):\n",
    "            with open(json_path, \"r\") as f:\n",
    "                self.infer_json_dict[job_name] = json.load(f)\n",
    "        \n",
    "        # check their length\n",
    "        first_key = next(iter(self.infer_json_dict))\n",
    "        for job_name, infer_json in self.infer_json_dict.items():\n",
    "            assert len(infer_json) == len(self.infer_json_dict[first_key]), f\"length not match for {job_name}\"\n",
    "        self._len = len(self.infer_json_dict[first_key])\n",
    "    \n",
    "    def __len__(self):\n",
    "        return self._len\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        return {job_name: infer_json[idx] for job_name, infer_json in self.infer_json_dict.items()}\n",
    "\n",
    "infer_json_dataset = MultiInferJson(infer_json_path_dict)\n",
    "\n",
    "def check_region_id_image_id(infer_json_dataset):\n",
    "    for sample in tqdm.tqdm(infer_json_dataset, desc=\"Check region_id and image_id\"):\n",
    "        first_key = next(iter(sample))\n",
    "        image_id = sample[first_key][\"metadata\"][\"metadata_image_id\"]\n",
    "        region_id = sample[first_key][\"metadata\"][\"metadata_region_id\"]\n",
    "        for job_name, region_pred in sample.items():\n",
    "            assert image_id == region_pred[\"metadata\"][\"metadata_image_id\"], f\"image_id not match for {job_name}\"\n",
    "            assert region_id == region_pred[\"metadata\"][\"metadata_region_id\"], f\"region_id not match for {job_name}\"\n",
    "\n",
    "check_region_id_image_id(infer_json_dataset)\n",
    "infer_json_dataset_iter = iter(infer_json_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_one_region(infer_json_dataset, region_cnt):\n",
    "    samples = infer_json_dataset[region_cnt]\n",
    "    first_key = next(iter(samples))\n",
    "    EVAL_DATASET_SPLIT = 'visual_genome-densecap-local-densecap-test'\n",
    "\n",
    "    first_sample = samples[first_key]\n",
    "\n",
    "    image_id = first_sample[\"metadata\"][\"metadata_image_id\"]\n",
    "    region_id = first_sample[\"metadata\"][\"metadata_region_id\"]\n",
    "    input_boxes = first_sample[\"metadata\"][\"metadata_input_boxes\"]\n",
    "\n",
    "    sample_cnt = image_id_to_dataset_id_mapping[EVAL_DATASET_SPLIT][str(image_id)]\n",
    "    sample = eval_dataset[EVAL_DATASET_SPLIT][sample_cnt]\n",
    "    image = sample[\"image\"]\n",
    "\n",
    "    references = first_sample[\"references\"]\n",
    "\n",
    "    candidates = []\n",
    "    for job_name, region_pred in samples.items():\n",
    "        candidates.extend(region_pred[\"candidates\"])\n",
    "\n",
    "    font_path = 'tmp/Arial.ttf'\n",
    "\n",
    "    # Calculate the number of colors  \n",
    "    num_colors = len(harmonious_colors)  \n",
    "    # Generate a random start index  \n",
    "    # start_index = random.randint(0, num_colors - 1)  \n",
    "    start_index = 4\n",
    "    # Select colors in a round-robin way  \n",
    "    selected_colors = [harmonious_colors[(start_index + i) % num_colors] for i in range(len(candidates))]  \n",
    "\n",
    "    captions_color = selected_colors + [(255,255,255)]\n",
    "    captions = candidates + references\n",
    "    \n",
    "    model_color_fig_path = os.path.join(training_args.output_dir, \"model_color_fig.png\")\n",
    "    if not os.path.exists(model_color_fig_path):\n",
    "        model_name = [job_name for job_name in samples.keys()]\n",
    "        model_color_fig = draw_captions(Image.new('RGB', (256, 0)), model_name, font_path, captions_color=selected_colors)\n",
    "        model_color_fig.save(model_color_fig_path)\n",
    "        print(f\"save model_color_fig to {model_color_fig_path}\")\n",
    "\n",
    "    pil_img_with_bbox_and_captions = plot_bbox_and_captions(image, input_boxes, captions, font_path, captions_color=captions_color, margin_size=5)  \n",
    "    return pil_img_with_bbox_and_captions, f\"{region_cnt}-{sample_cnt}-{region_id}-{image_id}.png\"\n",
    "\n",
    "region_cnt = 0\n",
    "pil_img_with_bbox_and_captions, pil_img_with_bbox_and_captions_path = plot_one_region(infer_json_dataset, region_cnt)\n",
    "pil_img_with_bbox_and_captions.save(os.path.join(training_args.output_dir, pil_img_with_bbox_and_captions_path))\n",
    "pil_img_with_bbox_and_captions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _add_prefix_suffix_to_path(path: str, prefix: str, suffix: str) -> str:\n",
    "    base_dir, filename = os.path.split(path)\n",
    "    return os.path.join(base_dir, prefix + filename + suffix)\n",
    "\n",
    "score_json_path_dict = {}\n",
    "# CIDEr-D-scores.infer-visual_genome-region_descriptions_v1.2.0-test.json.json\n",
    "SCORE_PREFIX = \"CIDEr-D-scores.\"\n",
    "SCORE_SUFFIX = \".json\"\n",
    "\n",
    "for k, v in infer_json_path_dict.items():\n",
    "    score_json_path_dict[k] = _add_prefix_suffix_to_path(v, SCORE_PREFIX, SCORE_SUFFIX)\n",
    "for job_name, score_json_path in score_json_path_dict.items():\n",
    "    print(f\"job_name: {job_name}\")\n",
    "    print(f\"is exists: {os.path.exists(score_json_path)}\")\n",
    "    assert os.path.exists(score_json_path), f\"{score_json_path} not exists\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import json\n",
    "\n",
    "score_json_dict = {}\n",
    "for k, v in score_json_path_dict.items():\n",
    "    with open(v, \"r\") as f:\n",
    "        score_json_dict[k] = json.load(f)\n",
    "def build_score_df(score_json_dict):\n",
    "    return pd.DataFrame.from_dict({k: v for k, v in score_json_dict.items()})\n",
    "score_df= build_score_df(score_json_dict)\n",
    "score_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot the dist of scores for different column\n",
    "import matplotlib.pyplot as plt\n",
    "def plot_score_desc(score_df):\n",
    "    fig, ax = plt.subplots(figsize=(5, 3))\n",
    "    for col_name in score_df.columns:\n",
    "        col_seris = score_df[col_name].sort_values(ascending=False)\n",
    "        col_values = col_seris.values\n",
    "        ax.plot(col_values, label=col_name)\n",
    "    ax.legend()\n",
    "    ax.set_xlabel(\"samples\")\n",
    "    ax.set_ylabel(\"score\")\n",
    "    return fig, ax\n",
    "fig, ax = plot_score_desc(score_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sorted_score_df = score_df.sort_values(by=\"o365_vg-gpt2l-bs_64-lsj\", ascending=False)\n",
    "sorted_score_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_regions = len(infer_json_dataset)\n",
    "\n",
    "sorted_score_seris = sorted_score_df.iloc[int(num_regions * 0.98521)]\n",
    "region_cnt = sorted_score_seris.name\n",
    "\n",
    "# region_cnt  = random.randint(0, num_regions-1)\n",
    "\n",
    "score_seris = score_df.iloc[region_cnt]\n",
    "pil_img_with_bbox_and_captions, pil_img_with_bbox_and_captions_path = plot_one_region(infer_json_dataset, region_cnt)\n",
    "# pil_img_with_bbox_and_captions.save(pil_img_with_bbox_and_captions_path)\n",
    "print(f\"region_cnt: {region_cnt}\\nscores: {score_seris}\\nsave to: {pil_img_with_bbox_and_captions_path}\")\n",
    "pil_img_with_bbox_and_captions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import base64  \n",
    "from PIL import Image  \n",
    "import io\n",
    "from IPython.display import display, HTML  \n",
    "  \n",
    "def visualize_images_html(infer_json_dataset, num_images, images_per_row=5):  \n",
    "    images_html = \"<table>\"  \n",
    "    region_cnt_random_list = np.random.randint(0, len(infer_json_dataset), num_images)\n",
    "    print(f\"The region cnt random list: {region_cnt_random_list}\")\n",
    "    for region_cnt in region_cnt_random_list:\n",
    "        if region_cnt % images_per_row == 0:  \n",
    "            images_html += \"<tr>\"\n",
    "        pil_img, pil_img_note = plot_one_region(infer_json_dataset, region_cnt)\n",
    "  \n",
    "        # Create an in-memory bytes buffer  \n",
    "        buf = io.BytesIO()  \n",
    "  \n",
    "        # Save the PIL image to the buffer in PNG format  \n",
    "        pil_img.save(buf, format='PNG')  \n",
    "  \n",
    "        # Get the base64 encoded string  \n",
    "        img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')  \n",
    "  \n",
    "        images_html += '<td><img src=\"data:image/png;base64,{}\"  height=\"500\"><br>{}</td>'.format(img_base64, pil_img_note)  \n",
    "        if region_cnt % images_per_row == images_per_row - 1:  \n",
    "            images_html += \"</tr>\"  \n",
    "  \n",
    "    images_html += \"</table>\"  \n",
    "    print(f\"html is ready!\")\n",
    "    display(HTML(images_html))  \n",
    "\n",
    "visualize_images_html(infer_json_dataset, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from flask import Flask, render_template  \n",
    "from PIL import Image  \n",
    "import io  \n",
    "import base64  \n",
    "  \n",
    "app = Flask(__name__)  \n",
    "  \n",
    "@app.route('/')  \n",
    "def home():  \n",
    "    num_images = 10  \n",
    "    images_per_row = 5  \n",
    "    images = []  \n",
    "    for i in range(num_images):  \n",
    "        pil_img, pil_img_note = plot_one_region(infer_json_dataset, region_cnt)  # Assuming dataset[i] returns a tuple of (image, caption)  \n",
    "        buf = io.BytesIO()  \n",
    "        pil_img.save(buf, format='PNG')  \n",
    "        img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')  \n",
    "        images.append((img_base64, pil_img_note))  \n",
    "    return render_template('tmp/home.html', images=images, images_per_row=images_per_row)  \n",
    "  \n",
    "if __name__ == '__main__':  \n",
    "    app.run(debug=True)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sca-v2",
   "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}