get_amazon_ocr.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "discrete-tolerance",
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "from botocore.client import Config\n",
    "import io\n",
    "import json\n",
    "import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "christian-import",
   "metadata": {},
   "outputs": [],
   "source": [
    "AWS_ACCESS_KEY_ID = ''\n",
    "AWS_SECRET_ACCESS_KEY = ''\n",
    "REGION_NAME = 'us-west-2'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "spectacular-touch",
   "metadata": {},
   "source": [
    "# TEXTVQA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "refined-exception",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = json.load(open('/media/abiten/4TB/Datasets/textvqa/train/TextVQA_0.5.1_train.json', 'r'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "distinguished-documentation",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 21953/21953 [9:10:55<00:00,  1.51s/it]  \n",
      "100%|██████████| 3166/3166 [1:21:26<00:00,  1.54s/it]\n",
      "100%|██████████| 3289/3289 [1:19:09<00:00,  1.44s/it]\n"
     ]
    }
   ],
   "source": [
    "splits = ['train', 'val', 'test']\n",
    "bucket = 'stvqa'\n",
    "\n",
    "for split in splits:\n",
    "    ocr_response = {}\n",
    "    data = json.load(open('/media/abiten/4TB/Datasets/textvqa/{}/TextVQA_0.5.1_{}.json'.format(split, split), 'r'))\n",
    "    img_ids = list(set([d['image_id'] for d in data['data']]))\n",
    "    config = Config(retries=dict(max_attempts=5))\n",
    "    rek_client = boto3.client('rekognition', aws_access_key_id=AWS_ACCESS_KEY_ID,\n",
    "                            aws_secret_access_key=AWS_SECRET_ACCESS_KEY,\n",
    "                            region_name=REGION_NAME,\n",
    "                           config=config)\n",
    "    if split == 'test':\n",
    "        aux_path = 'test_images'\n",
    "    else:\n",
    "        aux_path = 'train_images'\n",
    "    \n",
    "    for img_id in tqdm.tqdm(img_ids):\n",
    "\n",
    "        if split == 'val':\n",
    "            file_name = 'textvqa/train/{}/{}.jpg'.format(aux_path, img_id)\n",
    "        else:\n",
    "            file_name = 'textvqa/{}/{}/{}.jpg'.format(split, aux_path, img_id)\n",
    "\n",
    "        response = rek_client.detect_text(Image={'S3Object':{'Bucket':bucket,'Name':file_name}})\n",
    "        ocr_response[img_id] = response\n",
    "        \n",
    "    json.dump(ocr_response, open('/media/abiten/4TB/Datasets/textvqa/{}_amazon_ocr.json'.format(split), 'w'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "daily-evening",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "intensive-setting",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "pending-imperial",
   "metadata": {},
   "source": [
    "# STVQA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "agreed-confusion",
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "renewable-ireland",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2971/2971 [56:49<00:00,  1.15s/it] \n"
     ]
    }
   ],
   "source": [
    "splits = ['train', 'test']\n",
    "# splits = ['test']\n",
    "bucket = 'stvqa'\n",
    "config = Config(retries=dict(max_attempts=5))\n",
    "rek_client = boto3.client('rekognition', aws_access_key_id=AWS_ACCESS_KEY_ID,\n",
    "                        aws_secret_access_key=AWS_SECRET_ACCESS_KEY,\n",
    "                        region_name=REGION_NAME,\n",
    "                       config=config)\n",
    "for split in splits:\n",
    "    ocr_response = {}\n",
    "    data = json.load(open('/media/abiten/4TB/Datasets/st-vqa/{}_task_3.json'.format(split, split), 'r'))\n",
    "    img_ids = list(set([d['file_path'] for d in data['data']]))\n",
    "    config = Config(retries=dict(max_attempts=5))\n",
    "\n",
    "    for img_id in tqdm.tqdm(img_ids):\n",
    "        try:\n",
    "            if split == 'test':\n",
    "                file_name = 'st-vqa/{}/task3/{}'.format(split, img_id)\n",
    "            else:\n",
    "                file_name = 'st-vqa/{}/{}'.format(split, img_id)\n",
    "\n",
    "            response = rek_client.detect_text(Image={'S3Object':{'Bucket':bucket,'Name':file_name}})\n",
    "            ocr_response[img_id] = response\n",
    "        except:\n",
    "            print(img_id)\n",
    "    json.dump(ocr_response, open('/media/abiten/4TB/Datasets/st-vqa/{}_amazon_ocr.json'.format(split), 'w'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "knowing-willow",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9354"
      ]
     },
     "execution_count": 75,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(ocr_response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "egyptian-darwin",
   "metadata": {},
   "source": [
    "# We need to replace COCO-Text with the original size images otherwise Amazon OCR doesnt work well"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "acquired-vector",
   "metadata": {},
   "outputs": [],
   "source": [
    "coco_path = '/media/abiten/4TB/Datasets/COCO_raw/train2014'\n",
    "for split in splits:\n",
    "    data = json.load(open('/media/abiten/4TB/Datasets/st-vqa/{}_task_3.json'.format(split, split), 'r'))\n",
    "    coco_text = list(set([d['file_path'].split('/')[1] for d in data['data'] if 'coco-text' in d['file_path']]))\n",
    "    for ct_img in coco_text:\n",
    "        if split == 'test':\n",
    "            file_name = '{}/task3/coco-text/{}'.format(split, ct_img)\n",
    "        else:\n",
    "            file_name = '{}/coco-text/{}'.format(split, ct_img)\n",
    "        org_img = os.path.join(coco_path, ct_img)\n",
    "        shutil.copy(org_img, os.path.join('/media/abiten/4TB/Datasets/st-vqa/', file_name))\n",
    "#     orginal_img = '/media/abiten/4TB/Datasets/COCO_raw/train2014'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cheap-mediterranean",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "unavailable-fetish",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "laden-remainder",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:main] *",
   "language": "python",
   "name": "conda-env-main-py"
  },
  "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.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}