2-create-train-and-val-dataset.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "from pathlib import Path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_text_paths = sorted(list(Path('images').glob(\"*/*.txt\")))\n",
    "val_text_paths = [os.path.join(os.getcwd(), 'Test Images', img) for img in os.listdir(os.path.join(os.getcwd(), 'Test Images')) if img.endswith('.txt')]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_class(document_name):\n",
    "\n",
    "    ClassToBeReturned = -1\n",
    "    if document_name.lower().__contains__('w9'):\n",
    "        ClassToBeReturned = 1\n",
    "    elif document_name.lower().__contains__('iban'):\n",
    "        ClassToBeReturned = 2\n",
    "    elif document_name.lower().__contains__('iso'):\n",
    "        ClassToBeReturned = 3\n",
    "    \n",
    "    return ClassToBeReturned"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_df(text_paths): \n",
    "    document_name = []\n",
    "    document_context = []\n",
    "\n",
    "    for text_file in text_paths:\n",
    "        document_name.append(str(text_file).split(\"\\\\\")[-1])\n",
    "        with open(text_file, 'r') as f:\n",
    "            document_context.append(f.readlines()[0])\n",
    "\n",
    "    df = pd.DataFrame({\"Document Name\": document_name, 'Document Context': document_context})\n",
    "\n",
    "    df['Document Class'] = df['Document Name'].apply(lambda x: get_class(str(x)))\n",
    "\n",
    "    return df\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_df = create_df(train_text_paths)\n",
    "validation_df = create_df(val_text_paths)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2    38\n",
       "3    31\n",
       "1    30\n",
       "Name: Document Class, dtype: int64"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "training_df['Document Class'].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2    1\n",
       "3    1\n",
       "1    1\n",
       "Name: Document Class, dtype: int64"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "validation_df['Document Class'].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_df.to_csv('Train.csv', index=False)\n",
    "validation_df.to_csv('Validation.csv', index=False)"
   ]
  }
 ],
 "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.8.10"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "9650cb4e16cdd4a8e8e2d128bf38d875813998db22a3c986335f89e0cb4d7bb2"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}