From 0484ae01b1e0dec32702f1a4b0bf8a67abf354fb Mon Sep 17 00:00:00 2001
From: Falko Habel <falko.habel@gmx.de>
Date: Thu, 30 Jan 2025 13:21:56 +0100
Subject: [PATCH] added tqdm and removed doubles

---
 src/aiunn/finetune.py | 11 +++--------
 1 file changed, 3 insertions(+), 8 deletions(-)

diff --git a/src/aiunn/finetune.py b/src/aiunn/finetune.py
index 4bdf821..f26ef93 100644
--- a/src/aiunn/finetune.py
+++ b/src/aiunn/finetune.py
@@ -9,14 +9,9 @@ from aiia.model import AIIABase, AIIA
 from sklearn.model_selection import train_test_split
 from typing import Dict, List, Union, Optional
 import base64
+from tqdm import tqdm
 
 
-import torch
-from torch.utils.data import Dataset
-from PIL import Image, ImageFile
-import io
-import base64
-import pandas as pd
 
 class ImageDataset(Dataset):
     def __init__(self, dataframe, transform=None):
@@ -222,7 +217,7 @@ class ModelTrainer:
         """
         self.model.to(self.device)
 
-        for epoch in range(num_epochs):
+        for epoch in tqdm(num_epochs):
             print(f"Epoch {epoch+1}/{num_epochs}")
 
             # Train phase
@@ -243,7 +238,7 @@ class ModelTrainer:
         self.model.train()
         running_loss = 0.0
 
-        for batch in self.train_loader:
+        for batch in tqdm(self.train_loader):
             low_ress = batch['low_ress'].to(self.device)
             high_ress = batch['high_ress'].to(self.device)