Refactor dataset handling and enhance evaluation: update NatureFromFolder class for test dataset size, add evaluate_and_plot function in test.py, and improve config settings.

Author: halak0013

data/results.csv

@@ -2,11 +2,12 @@ experiment,time,psnr,ssim,loss_critic,loss_generator
 7_first_experiment,00:16:50,11.103890439714435,0.34202152,1.099349385527171,0.1359923942584425
 8_model_compile,00:16:49,19.3652867520302,0.4844975471496582,-1.2957792903326415,0.0911528160242729
 9_transform,00:17:01,1.7978309320343853,-0.0438352711498737,-7.376533726493926,0.7896081656897993
-10_normal,00:16:51,8.219669060046096,0.31684741377830505,-31.28607681788266,1.967616661362452
-14_multistep,04:07:37,23.004348053166062,0.5242958068847656,-2.5751562796301553,-0.035972834620931246
-15_cosine,04:02:15,21.85281842397048,0.5211549997329712,-1.9651767669255815,0.23522654886764408
-15_cosine_warm_restarts,04:06:25,22.68296964537476,0.5351459980010986,-1.6454665754569295,-0.03318920495223519
-17_step,04:06:31,14.040152634682375,0.2926912009716034,,
-18_gan_custom,04:05:33,23.52549501689562,0.5291714072227478,-2.1092974134947307,-0.043675970456390975
-19_exponential,04:13:36,22.58134529896099,0.5127521753311157,-1.2576447154311854,0.07773792746407023
-20_lambda,04:07:58,22.653632716710888,0.5084766745567322,-1.1908818194046753,0.12637249500213202
+10_normal,00:16:51,8.219669060046096,0.316847413778305,-31.28607681788266,1.967616661362452
+14_multistep,04:07:37,23.004348053166066,0.5242958068847656,-2.5751562796301557,-0.0359728346209312
+15_cosine,04:02:15,21.85281842397048,0.5211549997329712,-1.9651767669255813,0.235226548867644
+15_cosine_warm_restarts,04:06:25,22.68296964537476,0.5351459980010986,-1.6454665754569295,-0.0331892049522351
+17_step,04:06:31,14.040152634682377,0.2926912009716034,,
+18_gan_custom,04:05:33,23.52549501689562,0.5291714072227478,-2.1092974134947307,-0.0436759704563909
+19_exponential,04:13:36,22.58134529896099,0.5127521753311157,-1.2576447154311854,0.0777379274640702
+20_lambda,04:07:58,22.653632716710888,0.5084766745567322,-1.1908818194046753,0.126372495002132
+21_Custom_new_dataset,12:40:16,30.987861233820492,0.6934245824813843,-0.8939516761237561,0.09604286165589182

src/test.py

@@ -1,19 +1,94 @@
 import torch
 from torch import optim
+import matplotlib.pyplot as plt
+import numpy as np
+from skimage.metrics import peak_signal_noise_ratio, structural_similarity
+from tqdm import tqdm
+
 
 from src.esrgan.model import Generator
 from src.utils import config
 from src.utils.utils import load_checkpoint, plot_examples
 from src.utils.utils import seed_torch
+from src.utils.data_loaders import get_loaders
+
+import matplotlib.pyplot as plt
+import numpy as np
+from skimage.metrics import peak_signal_noise_ratio, structural_similarity
+
+
+def evaluate_and_plot(test_loader, gen, num_samples=3):
+    psnr_list = []
+    ssim_list = []
+    selected_samples = []
+
+    loop = tqdm(test_loader, desc="Evaluating", leave=False)
+    with torch.no_grad():
+        for idx, (low_res, high_res) in enumerate(loop):
+            low_res = low_res.to(config.DEVICE)
+            high_res = high_res.to(config.DEVICE)
+
+            # Generator ile yüksek çözünürlüklü görüntü üret
+            fake_high_res = gen(low_res)
+
+            # PSNR ve SSIM hesapla
+            for i in range(low_res.shape[0]):
+                sr_img = fake_high_res[i].cpu().numpy().transpose(1, 2, 0)
+                hr_img = high_res[i].cpu().numpy().transpose(1, 2, 0)
+                lr_img = low_res[i].cpu().numpy().transpose(1, 2, 0)
+
+                # Normalize görüntüler (0-1 aralığına)
+                sr_img = np.clip(sr_img, 0, 1)
+                hr_img = np.clip(hr_img, 0, 1)
+                lr_img = np.clip(lr_img, 0, 1)
+
+                # PSNR ve SSIM hesapla
+                psnr = peak_signal_noise_ratio(hr_img, sr_img, data_range=1.0)
+                ssim = structural_similarity(
+                    hr_img, sr_img, channel_axis=2, data_range=1.0
+                )
+                psnr_list.append(psnr)
+                ssim_list.append(ssim)
+
+                # İlk num_samples kadar örneği seç
+                if len(selected_samples) < num_samples:
+                    selected_samples.append((lr_img, hr_img, sr_img, psnr, ssim))
+
+    # Ortalama PSNR ve SSIM değerlerini yazdır
+    print(f"Average PSNR: {np.mean(psnr_list):.4f}")
+    print(f"Average SSIM: {np.mean(ssim_list):.4f}")
+
+    # Seçilen örnekleri çiz
+    fig, axes = plt.subplots(num_samples, 3, figsize=(10, 3 * num_samples))
+    fig.suptitle(
+        f"Super-Resolution Evaluation Samples\nPSNR: {np.mean(psnr_list):.4f} SSIM{np.mean(ssim_list):.4f}",
+        fontsize=16,
+    )
+    for i, (lr, hr, sr, psnr, ssim) in enumerate(selected_samples):
+        axes[i, 0].imshow(lr)
+        axes[i, 0].set_title(f"LR")
+        axes[i, 0].axis("off")
+
+        axes[i, 1].imshow(sr)
+        axes[i, 1].set_title(f"SR  PSNR: {psnr:.2f} SSIM: {ssim:.4f}")
+        axes[i, 1].axis("off")
+
+        axes[i, 2].imshow(hr)
+        axes[i, 2].set_title(f"HR")
+        axes[i, 2].axis("off")
+
+    plt.tight_layout()
+    plt.savefig(
+        f"{config.SAVE_PATH}/evaluation_samples.png", dpi=300, bbox_inches="tight"
+    )
+    plt.show()
+
 
 def test():
     seed_torch(config.SEED)
+    _, _, test_loader = get_loaders()
     gen = Generator(in_channels=3).to(config.DEVICE)
-    opt_gen = optim.Adam(gen.parameters(), lr=config.LEARNING_RATE, betas=(0.0, 0.9))
-    load_checkpoint(
-        config.CHECKPOINT_GEN,
-        gen,
-        opt_gen,
-        config.LEARNING_RATE,
-    )
-    plot_examples(config.TEST_IMAGE_DIR, gen)
+    gen.eval()
+    load_checkpoint(config.CHECKPOINT_GEN, gen)
+    print("Testing the model...")
+    evaluate_and_plot(test_loader, gen, num_samples=12)

src/utils/config.py

@@ -13,6 +13,7 @@
 LAMBDA_GP = 10
 NUM_WORKERS = 12
 NUM_EPOCHS = 25
+TEST_SIZE = 1000
 
 BATCH_SIZE = 32
 HIGH_RES = 128
@@ -30,8 +31,12 @@
 LOG_DIR = "logs"
 DATA_SET_NAME_DIR = "esrgan_dataset"
 EXPERIMENT = VERSION + "_" + NAME
+
+
 def extension(mode_name: str) -> str:
     return mode_name + EXPERIMENT + VERSION + ".pth"
+
+
 SAVE_PATH = f"{LOG_DIR}/{DATA_SET_NAME_DIR}/{EXPERIMENT}"
 DATA = "data"
 CACHE_DIR = DATA + "/cache"

src/utils/dataset_from_folder.py

@@ -15,7 +15,7 @@ def __init__(self, type_="train", train_ratio=0.7, val_ratio=0.1):
         self.class_names = os.listdir(self.root_dir)
         print(self.class_names[:10])
 
-        for index, name in enumerate(self.class_names):
+        for name in self.class_names:
             files = os.path.join(self.root_dir, name)
             self.data.append(files)
         print(self.data[:10])
@@ -32,7 +32,10 @@ def __init__(self, type_="train", train_ratio=0.7, val_ratio=0.1):
         elif type_ == "val":
             self.data = self.data[train_end:val_end]
         elif type_ == "test":
-            self.data = self.data[val_end:]
+            test_size = total - val_end
+            if test_size >= cfg.TEST_SIZE:
+                test_size = cfg.TEST_SIZE
+            self.data = self.data[val_end : val_end + test_size]
         else:
             raise ValueError("type_ must be 'train', 'val', or 'test'")
 
@@ -67,4 +70,4 @@ def test():
 
 
 if __name__ == "__main__":
-    test()
+    test()

src/utils/utils.py

@@ -27,7 +27,7 @@ def gradient_penalty(discriminator, real, fake):
         retain_graph=True,
     )[0]
     # Her görüntüyü vektörleştirir
-    gradient = gradient.flatten(start_dim=1) 
+    gradient = gradient.flatten(start_dim=1)
     # Her örneğin gradyan uzunluğunu hesaplar
     gradient_norm = gradient.norm(2, dim=1)
     # Normun 1’den sapmasına ceza verir
@@ -44,17 +44,18 @@ def save_checkpoint(model, optimizer, filename="my_checkpoint.pth.tar"):
     torch.save(checkpoint, filename)
 
 
-def load_checkpoint(checkpoint_file, model, optimizer, lr):
+def load_checkpoint(checkpoint_file, model, optimizer=None, lr=None):
     print("=> Loading checkpoint")
     checkpoint = torch.load(checkpoint_file, map_location=cfg.DEVICE)
     # model.load_state_dict(checkpoint)
     model.load_state_dict(checkpoint["state_dict"])
-    optimizer.load_state_dict(checkpoint["optimizer"])
 
     # If we don't do this then it will just have learning rate of old checkpoint
     # and it will lead to many hours of debugging \:
-    for param_group in optimizer.param_groups:
-        param_group["lr"] = lr
+    if lr is not None:
+        optimizer.load_state_dict(checkpoint["optimizer"])
+        for param_group in optimizer.param_groups:
+            param_group["lr"] = lr
 
 
 def plot_examples(low_res_folder, gen, ex):