From e4c624e521f00c03fdb5db4379cf89959597c915 Mon Sep 17 00:00:00 2001
From: Abu Bakr <abubakr.soliman@alefeducation.com>
Date: Fri, 17 Dec 2021 06:40:07 +0200
Subject: [PATCH] support training progress bar

---
 README.md | 2 +-
 train.py  | 4 ++++
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 50de47c4e0..64c49a4b31 100644
--- a/README.md
+++ b/README.md
@@ -27,7 +27,7 @@ The difference between our paper and ICDAR challenge is summarized [here](https:
 In the paper, expriments were performed with **PyTorch 0.4.1, CUDA 9.0**.
 - requirements : lmdb, pillow, torchvision, nltk, natsort
 ```
-pip3 install lmdb pillow torchvision nltk natsort
+pip3 install lmdb pillow torchvision nltk natsort tqdm
 ```
 
 ### Download lmdb dataset for traininig and evaluation from [here](https://www.dropbox.com/sh/i39abvnefllx2si/AAAbAYRvxzRp3cIE5HzqUw3ra?dl=0)
diff --git a/train.py b/train.py
index 7f3a66bf00..5f3dd8e89e 100644
--- a/train.py
+++ b/train.py
@@ -11,6 +11,7 @@
 import torch.optim as optim
 import torch.utils.data
 import numpy as np
+from tqdm.auto import tqdm
 
 from utils import CTCLabelConverter, CTCLabelConverterForBaiduWarpctc, AttnLabelConverter, Averager
 from dataset import hierarchical_dataset, AlignCollate, Batch_Balanced_Dataset
@@ -142,6 +143,7 @@ def train(opt):
     best_norm_ED = -1
     iteration = start_iter
 
+    pbar = tqdm(total=opt.num_iter)
     while(True):
         # train part
         image_tensors, labels = train_dataset.get_batch()
@@ -224,6 +226,8 @@ def train(opt):
             print('end the training')
             sys.exit()
         iteration += 1
+        pbar.update(1)
+    pbar.close()
 
 
 if __name__ == '__main__':