From 07e8563c5d803e86055e37acbe532e686882ca67 Mon Sep 17 00:00:00 2001
From: Tinsae <tinsaekahsay@gmail.com>
Date: Fri, 1 May 2026 23:36:25 +0200
Subject: [PATCH] separated train and test transformation

---
 Data.py           | 29 +++++++++++++++++++++++------
 IdentitySubset.py | 13 ++++++++++---
 Tune.py           | 35 ++++++++++++++++++++++-------------
 lib64             |  1 +
 4 files changed, 56 insertions(+), 22 deletions(-)
 create mode 120000 lib64

diff --git a/Data.py b/Data.py
index 3e20860..53eaf95 100644
--- a/Data.py
+++ b/Data.py
@@ -2,12 +2,18 @@ from torchvision import datasets, transforms, models
 import torch
 import numpy as np
 
-# transform images to size
-def transform(res):
+# train set transform 
+def train_transform(res):
     return transforms.Compose([
         # ResNet expects 224 x 224 res 
+        # Inception expects 299 x 299
         transforms.Resize((res, res)),
-        transforms.RandomHorizontalFlip(),
+        transforms.RandomHorizontalFlip(p=0.5),
+        transforms.ColorJitter(
+            brightness=0.2, 
+            contrast=0.2, 
+            saturation=0.1
+        ),
         transforms.ToTensor(),
         # normalise to 
         transforms.Normalize(
@@ -16,14 +22,26 @@ def transform(res):
         )
     ])
 
+# test set transform
+def test_transform(res):
+    return transforms.Compose([
+        # Just standard resize to 224x224
+        transforms.Resize((res, res)),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406], 
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+
 # Load data with 'identity' as target and transform it
-def get_set(res):
+def get_set():
     return datasets.CelebA(
         root='./data', 
         split='all', 
         target_type='identity', 
         download=True, 
-        transform=transform(res)
+        transform=None
     )
 
 
@@ -66,7 +84,6 @@ def get_indices(dataset, identities, split_at):
     test_indices = []
 
     #training_sample = int(sample_size * training_ratio)
-
     for person_id in identities:
         # Get all indices for this specific person
         indices = torch.where(dataset.identity == person_id)[0].numpy()
diff --git a/IdentitySubset.py b/IdentitySubset.py
index 655b393..5db7685 100644
--- a/IdentitySubset.py
+++ b/IdentitySubset.py
@@ -2,12 +2,19 @@
 import torch
 
 class IdentitySubset(torch.utils.data.Dataset):
-    def __init__(self, full_ds, indices, id_mapping):
-        self.full_ds = full_ds
+    def __init__(self, dataset, indices, id_mapping, transform=None):
+        self.dataset = dataset
         self.indices = indices
         self.id_mapping = id_mapping
+        self.transform = transform
+
     def __getitem__(self, idx):
-        img, old_id = self.full_ds[self.indices[idx]]
+        img, old_id = self.dataset[self.indices[idx]]
+
+        if self.transform:
+            img = self.transform(img)
+
         return img, self.id_mapping[old_id.item()]
+    
     def __len__(self):
         return len(self.indices)
\ No newline at end of file
diff --git a/Tune.py b/Tune.py
index 52d670a..845bfd4 100644
--- a/Tune.py
+++ b/Tune.py
@@ -8,15 +8,15 @@ from IdentitySubset import IdentitySubset
 from architectures.Model import Model, Architecture
 
 # numbre of classes 
-CLASS_SIZE = 30
+CLASS_SIZE = 20
 # batch 
-BATCH_SIZE = 16
+BATCH_SIZE = 8
 
 # size of images per class trainset + testset
 # 30 works best, more than that and we dont have enough data
 SAMPLE_SIZE = 30
 
-# this is then full sample - test sample
+# this is then (full_sample - test_sample)
 TRAINING_SMPLE = 28
 
 # learning rate
@@ -26,14 +26,18 @@ EPOCHS = 20
 # depends on model architecture
 # ResNet, DenseNet = 224
 # Inception = 299
-RESOLUTION = 299
+RESOLUTION = 224
 
-# model architecture
-arch = Architecture.INCEPTION
+# model architecture options are 
+#  - RESNET18
+#  - RESNET50
+#  - DENSENET121
+#  - INCEPTION
+arch = Architecture.RESNET18
 
 # DATA PREPARATION
 # load data set and prepare
-dataset = get_set(res = RESOLUTION)
+dataset = get_set()
 # select identities for experiment
 selected_identities = select_ids(
     dataset = dataset,
@@ -56,10 +60,12 @@ id_map = {old_id: new_id for new_id, old_id in enumerate(selected_identities)}
 
 # we remap identities because crossEntropyLoss requires in indices 0 -> (n-1)
 # where n = class size.
+tr_transform = train_transform(res = RESOLUTION)
 train_data = IdentitySubset(
-    dataset, 
-    train_indices, 
-    id_map)
+    dataset=dataset, 
+    indices=train_indices, 
+    id_mapping=id_map,
+    transform=tr_transform)
 
 train_loader = DataLoader(
     train_data, 
@@ -94,11 +100,14 @@ torch.save(
 print('Model saved!')
 
 # EVALUATE
+
+te_transform = test_transform(RESOLUTION)
 # Testing
 test_data = IdentitySubset(
-    dataset, 
-    test_indices, 
-    id_map)
+    dataset = dataset, 
+    indices=test_indices, 
+    id_mapping=id_map,
+    transform=te_transform)
 
 test_loader = DataLoader(
     test_data, 
diff --git a/lib64 b/lib64
new file mode 120000
index 0000000..7951405
--- /dev/null
+++ b/lib64
@@ -0,0 +1 @@
+lib
\ No newline at end of file