generate_synth_data.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import h5py\n",
    "import os, sys, traceback\n",
    "from synthgen import *\n",
    "from common import *\n",
    "import pickle\n",
    "from glob import glob\n",
    "import os\n",
    "import math\n",
    "from PIL import Image\n",
    "from multiprocessing import Pool\n",
    "from tqdm import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "import warnings\n",
    "warnings.simplefilter('ignore')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "В цьому ноутбуці ви знадете код для генерації синтетичних картинок з текстом, як описано в статті [\"Synthetic Data for Text Localisation in Natural Images\"](http://www.robots.ox.ac.uk/~vgg/data/scenetext/) з деякими модифікаціями: доданий код для генерації обрізаних текстових bounding boxes спеціально для задачі text recognition, так званих кропів. Також додана можливість генерувати не лише англійський текст, а і будь-якою іншою мовою. Для цього просто створіть свій текстовий файл, помістьть туди який побажаєте корпус і вкажіть шлях до файлу в параметр `TEXT_CORPUS_FILE`. Ми рекомендуємо зробити певний препроцесінг для тексту: видалити емоджі, \"погані символи\" і т д.\n",
    "Також майже все тут зберігається в h5 файлах. Що це і як з ними правильно працювати можна почитати [тут](http://docs.h5py.org/en/stable/build.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_IMG = -1 #number of images to use for generation (-1 to use all available):\n",
    "INSTANCE_PER_IMAGE = 10 #number of times to use the same image\n",
    "SECS_PER_IMG = 5 # max time per image in seconds\n",
    "TEXT_CORPUS_FILE = 'data/newsgroup/newsgroup.txt' # file with text corpus to use in generation\n",
    "PATH_TO_IMNAMES = 'imnames.cp' #file with background image names\n",
    "IMAGES_DIRECTORY = 'bg_img/' #directory with background images \n",
    "PATH_TO_DEPTH_DATABASE = 'depth.h5' #file with depth for each background image\n",
    "PATH_TO_SEG_DATABASE = 'seg.h5' #file with segmentation for each background image\n",
    "SAVE_IMAGES = False #set True if you want to store not only crops, but also whole images for text detection task (or end-to-end)\n",
    "SYNTH_IMS_DIR = 'results/' #where to store synthetic images if SAVE_IMAGES \n",
    "prefix_for_results = 'synth_text_result_' #this prefix will be used for each batch of images if SAVE_IMAGES\n",
    "SAVE_CROPS = True # set True if you want to store crops for text recognition task\n",
    "CROPS_DIR = 'crops_sample/' #where to store crops if SAVE_CROPS \n",
    "viz = False #set True if you want to visualize each synththetic image\n",
    "\n",
    "\n",
    "num_threads = 10 #\n",
    "batch_size = 50 # how much background images use to store data in one h5 base (for backup)\n",
    "\n",
    "\n",
    "if SAVE_CROPS and not os.path.exists(CROPS_DIR):\n",
    "    os.mkdir(CROPS_DIR)\n",
    "if SAVE_IMAGES and not os.path.exists(SYNTH_IMS_DIR):\n",
    "    os.mkdir(SYNTH_IMS_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "depth_db = h5py.File(PATH_TO_DEPTH_DATABASE,'r')\n",
    "seg_db = h5py.File(PATH_TO_SEG_DATABASE, 'r')['mask']\n",
    "imnames = pickle.load(open(PATH_TO_IMNAMES, 'rb'))\n",
    "images = [os.path.basename(x) for x in glob(f\"{IMAGES_DIRECTORY}/*jpg\")]\n",
    "imnames = list(filter(lambda x: x in images, imnames))  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_res_to_db(imgname,res,db):\n",
    "    \"\"\"\n",
    "    Add the synthetically generated text image instance\n",
    "    and other metadata to the dataset.\n",
    "    \"\"\"\n",
    "    ninstance = len(res)\n",
    "    for i in range(ninstance):\n",
    "        dname = \"%s_%d\"%(imgname, i)\n",
    "        db['data'].create_dataset(dname,data=res[i]['img'])\n",
    "        db['data'][dname].attrs['charBB'] = res[i]['charBB']\n",
    "        db['data'][dname].attrs['wordBB'] = res[i]['wordBB']  \n",
    "        L = res[i]['txt']\n",
    "        L = [n.encode(\"utf-8\") for n in L]\n",
    "        db['data'][dname].attrs['txt'] = L"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Це код для того, щоб отримати рівний, правильно повернутий кроп з bounding box. Інколи він може видавати помилки(наприклад, коли текст занадто близько до краю), а іноді - неправильно повертати текст і віддавати перевернутий. В продакшині в нас працює певний воркераунд: ми віддаємо дві картинки, ту, яку віддав наш метод і повернуту на 180 градусів, але це не надто хороше рішення, якщо ми хочемо обробляти набагато більший потік картинок. Тому feel free до будь-яких модифікацій."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_crop(pil_image, poly):\n",
    "    rotate_180 = None\n",
    "    img = np.array(pil_image)[:, :, ::-1]\n",
    "    polygon = [poly['x0'], poly['y0'], poly['x1'], poly['y1'], poly['x2'], poly['y2'], poly['x3'], poly['y3']]\n",
    "    height = math.sqrt((polygon[6] - polygon[0]) ** 2 + (polygon[7] - polygon[1]) ** 2)\n",
    "    width = math.sqrt((polygon[2] - polygon[0]) ** 2 + (polygon[3] - polygon[1]) ** 2)\n",
    "    try:\n",
    "        angle = math.pi - math.asin((polygon[7] - polygon[5]) / width)\n",
    "        rotate_180 = True\n",
    "    except:\n",
    "        cr = pil_image.crop((min(polygon[0], polygon[6]),\n",
    "                             min(polygon[1], polygon[3]),\n",
    "                             max(polygon[4], polygon[2]),\n",
    "                             max(polygon[5], polygon[7]))).convert('RGB')\n",
    "        try:\n",
    "            z = np.array(cr)[:, :, ::-1]\n",
    "        except:\n",
    "            return None\n",
    "    else:\n",
    "        x_center = (polygon[0] + polygon[4]) / 2\n",
    "        y_center = (polygon[1] + polygon[5]) / 2\n",
    "        rect = ((x_center, y_center), (width, height), angle * 180 / math.pi)\n",
    "        z = crop_minAreaRect(img, rect)\n",
    "        if 0 in z.shape:\n",
    "            angle = math.asin((polygon[3] - polygon[5]) / height) + math.pi / 2\n",
    "            rect = ((x_center, y_center), (width, height), angle * 180 / math.pi)\n",
    "            z = crop_minAreaRect(img, rect)\n",
    "            rotate_180 = False\n",
    "    if z.shape[0] / z.shape[1] > 0.8:\n",
    "        z = np.rot90(z, 3)\n",
    "    if rotate_180:\n",
    "        z = np.rot90(z, 2)\n",
    "    return z\n",
    "\n",
    "\n",
    "def crop_minAreaRect(img, rect):\n",
    "    # rotate img\n",
    "    angle = rect[2]\n",
    "    rows, cols = img.shape[0], img.shape[1]\n",
    "    M = cv2.getRotationMatrix2D((cols / 2, rows / 2), angle, 1)\n",
    "\n",
    "    c = round((rows ** 2 + cols ** 2) ** 0.5)\n",
    "    img_rot = cv2.warpAffine(img, M, (c, c))\n",
    "    box = cv2.boxPoints(rect)\n",
    "    pts = np.int0(cv2.transform(np.array([box]), M))[0]\n",
    "    pts[pts < 0] = 0\n",
    "    # crop\n",
    "    img_crop = img_rot[pts[1][1]:pts[0][1],\n",
    "               pts[1][0]:pts[2][0]]\n",
    "\n",
    "    return img_crop"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "А тут знаходиться код для самого вирізання кропів і зберігання розмітки. Як ви можете побачити, ми встановили обмеження на мінімальну довжину тексту 4. Ви можете прибрати чи змінити його, додати нові або не використовувати обмежень взагалі і подати дужееее правильний корпус"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_crops(res):\n",
    "    num_images_exists = len(glob(f\"{CROPS_DIR}/*jpg\")) \n",
    "    ninstance = len(res)\n",
    "    for k in range(ninstance):\n",
    "        rgb = res[k]['img']\n",
    "        wordBB = res[k]['wordBB']\n",
    "        txt = res[k]['txt']\n",
    "\n",
    "        res_txt = []\n",
    "        for word in txt:\n",
    "            if len(word.split()) > 1:\n",
    "                res_txt.extend(word.split())\n",
    "            else:\n",
    "                res_txt.append(word)\n",
    "        txt = res_txt\n",
    "        image = Image.fromarray(rgb)\n",
    "        \n",
    "        for i in range(wordBB.shape[-1]):\n",
    "            bb = wordBB[:,:,i]\n",
    "            bb = np.c_[bb,bb[:,0]]\n",
    "            image = Image.fromarray(rgb)\n",
    "            poly = {\n",
    "            'x0': bb[0, 0],\n",
    "            'x1': bb[0, 1],\n",
    "            'x2': bb[0, 2],\n",
    "            'x3': bb[0, 3],\n",
    "            'y0': bb[1, 0],\n",
    "            'y1': bb[1, 1],\n",
    "            'y2': bb[1, 2],\n",
    "            'y3': bb[1, 3],\n",
    "            }\n",
    "            label = txt[i]\n",
    "            try:\n",
    "                img_cropped = get_crop(image, poly)\n",
    "                crop_pil = Image.fromarray(img_cropped)\n",
    "            except Exception as e:\n",
    "                print(e)\n",
    "                continue\n",
    "            if len(label) >=4:\n",
    "                name = f\"{CROPS_DIR}/{num_images_exists}.jpg\"\n",
    "                crop_pil.save(name)\n",
    "                with open(name.replace('jpg', 'txt'), 'w', encoding='utf-8') as annotation:\n",
    "                    annotation.write(label + '\\n')\n",
    "                num_images_exists += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def viz_textbb(text_im, charBB_list, wordBB, alpha=1.0):\n",
    "    \"\"\"\n",
    "    text_im : image containing text\n",
    "    charBB_list : list of 2x4xn_i bounding-box matrices\n",
    "    wordBB : 2x4xm matrix of word coordinates\n",
    "    \"\"\"\n",
    "    plt.figure(figsize=(20, 20))\n",
    "    plt.imshow(text_im)\n",
    "\n",
    "    H,W = text_im.shape[:2]\n",
    "\n",
    "    # plot the character-BB:\n",
    "    for i in range(len(charBB_list)):\n",
    "        bbs = charBB_list[i]\n",
    "        ni = bbs.shape[-1]\n",
    "        for j in range(ni):\n",
    "            bb = bbs[:,:,j]\n",
    "            bb = np.c_[bb,bb[:,0]]\n",
    "            plt.plot(bb[0,:], bb[1,:], 'r', alpha=alpha/2)\n",
    "\n",
    "    # plot the word-BB:\n",
    "    for i in range(wordBB.shape[-1]):\n",
    "        bb = wordBB[:,:,i]\n",
    "        bb = np.c_[bb,bb[:,0]]\n",
    "        plt.plot(bb[0,:], bb[1,:], 'g', alpha=alpha)\n",
    "        # visualize the indiv vertices:\n",
    "        vcol = ['r','g','b','k']\n",
    "        for j in range(4):\n",
    "            plt.scatter(bb[0,j],bb[1,j],color=vcol[j])        \n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "global_batch = num_threads * batch_size\n",
    "\n",
    "\n",
    "def preprocess_batch(name_and_batch):\n",
    "    name_db, batch = name_and_batch\n",
    "    if SAVE_IMAGES:\n",
    "        out_db = h5py.File(name_db,'w')\n",
    "        out_db.create_group('/data')\n",
    "    for imname in batch:\n",
    "        try:\n",
    "            img = Image.open(os.path.join(IMAGES_DIRECTORY, imname))\n",
    "            \n",
    "            depth = depth_db[imname][:].T\n",
    "            depth = depth[:,:,1]\n",
    "            seg = seg_db[imname][:].astype('float32')\n",
    "            area = seg_db[imname].attrs['area']\n",
    "            label = seg_db[imname].attrs['label']\n",
    "\n",
    "            sz = depth.shape[:2][::-1]\n",
    "            img = np.array(img.resize(sz,Image.ANTIALIAS))\n",
    "            seg = np.array(Image.fromarray(seg).resize(sz,Image.NEAREST))\n",
    "            res = RV3.render_text(img,depth,seg,area,label,\n",
    "                                ninstance=INSTANCE_PER_IMAGE,viz=False)\n",
    "            if viz:\n",
    "                 for k in range(len(res)):\n",
    "                    rgb = res[k]['img']\n",
    "                    charBB = res[k]['charBB']\n",
    "                    wordBB = res[k]['wordBB']\n",
    "                    txt = res[k]['txt']\n",
    "                    viz_textbb(rgb, [charBB], wordBB)\n",
    "            if len(res) > 0:\n",
    "                if SAVE_CROPS:\n",
    "                    generate_crops(res)\n",
    "                if SAVE_IMAGES:\n",
    "                    add_res_to_db(imname,res,out_db)\n",
    "        except Exception as e:\n",
    "            continue\n",
    "    if SAVE_IMAGES:\n",
    "        out_db.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "А це власне код для генерації картинок. Це може зайняти певний час - кілька годин"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = len(imnames)\n",
    "if NUM_IMG < 0:\n",
    "    NUM_IMG = N\n",
    "RV3 = RendererV3('./data/',TEXT_CORPUS_FILE, max_time=SECS_PER_IMG)\n",
    "\n",
    "for i in tqdm(range(min(NUM_IMG, N) // global_batch)):\n",
    "    names_global_batch = imnames[i * global_batch: (i + 1) * global_batch]\n",
    "    batches = []\n",
    "    names = []\n",
    "    for j in range(num_threads):\n",
    "        batches.append(imnames[j * batch_size: (j + 1) * batch_size])\n",
    "        names.append(f'{SYNTH_IMS_DIR}/{prefix_for_results}{i}_{j}.h5')\n",
    "    list_to_preprocess = list(zip(names, batches))\n",
    "    pool = Pool(num_threads)\n",
    "    pool.map(preprocess_batch, list_to_preprocess)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}