line-in-memory.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.transforms import Compose\n",
    "\n",
    "from htr_crnn_ctc.datasource import ParquetDataSource, InMemoryDataSource\n",
    "from htr_crnn_ctc.dataset import CTCDataset\n",
    "from htr_crnn_ctc.transforms import Deslant, Rescale, ToRGB, ToTensor, Normalise\n",
    "from htr_crnn_ctc.dataloader import CTCDataLoader\n",
    "from htr_crnn_ctc.model import CTCModel\n",
    "from htr_crnn_ctc.learn import Learner\n",
    "from htr_crnn_ctc.utils import get_decode_map\n",
    "\n",
    "from torch import cuda, device as Device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dev = Device(\"cuda\" if cuda.is_available() else \"cpu\")\n",
    "print(f\"dev: {dev}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pds = ParquetDataSource(\n",
    "#     file=\"tmp\\\\dataset\\\\IAM-line\\\\data\\\\train.parquet\",\n",
    "#     map_columns=None\n",
    "# )\n",
    "\n",
    "imds = InMemoryDataSource(file_path=\"tmp\\\\dataset\\\\IAM-line\\\\data\\\\train.parquet.deslanted.bin\")\n",
    "# imds.from_datasource(pds)\n",
    "# imds.dump()\n",
    "imds.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pre_trans = [\n",
    "    Deslant()\n",
    "]\n",
    "\n",
    "trans = [\n",
    "    Rescale(\n",
    "        output_size=(64, 800),\n",
    "        random_pad=True,\n",
    "        border_pad=(10, 40), \n",
    "        random_rotation=2,\n",
    "        random_stretch=1.2\n",
    "    ),\n",
    "    ToRGB(), \n",
    "    ToTensor(rgb=True),\n",
    "    Normalise(\n",
    "        mean=[0.485, 0.456, 0.406],\n",
    "        std=[0.229, 0.224, 0.225]\n",
    "    )\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = CTCDataset(\n",
    "    data_source=imds, # pds / imds\n",
    "    char_dict=None,\n",
    "    transform=Compose(trans) # pre_trans / trans\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for i in range(len(ds)):\n",
    "#     imds[i] = ds[i]\n",
    "\n",
    "# imds.dump()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dl = CTCDataLoader(\n",
    "    dataset=ds,\n",
    "    train_batch_size=120,\n",
    "    validation_batch_size=240,\n",
    "    validation_split=0.2,\n",
    "    shuffle=True,\n",
    "    seed=42,\n",
    "    device=dev\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = CTCModel(\n",
    "    chan_in=3,\n",
    "    time_step=96,\n",
    "    feature_size=512,\n",
    "    hidden_size=512,\n",
    "    output_size=len(ds.char_dict) + 1,\n",
    "    num_rnn_layers=4,\n",
    "    rnn_dropout=0\n",
    ").to(dev)\n",
    "\n",
    "model.load_pretrained_resnet()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn = Learner(\n",
    "    model=model,\n",
    "    dataloader=dl,\n",
    "    decode_map=get_decode_map(ds.char_dict)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# learn.freeze()\n",
    "# log, lr = learn.find_lr(start_lr=1e-5, end_lr=1e1, wd=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.freeze()\n",
    "learn.fit_one_cycle(epochs=5, max_lr=1e-3, base_lr=1e-4, wd=0.1)\n",
    "\n",
    "learn.unfreeze()\n",
    "learn.fit_one_cycle(epochs=5, max_lr=1e-3, base_lr=1e-4, wd=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# learn.predict(trans, img_ndarray=imds[0].image, dev=Device(\"cuda\"))\n",
    "learn.predict(Compose(pre_trans + trans), img_path=\"test.png\", dev=Device(\"cuda\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "learn.batch_predict(show_img=True, up_to=10)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}