added train,visualize_results scripts

configs/experiment1.yaml

@@ -6,9 +6,11 @@ model: model_v1
 # common parameters
 seed: 324
 train_size: 0.8
-image_size: 224
+# image_size: 224
+image_size: 64
 shuffle: False
 
 # training related
 batch_size: 16
-epochs: 10
+# epochs: 10
+epochs: 02

src/scripts/{create_sub_task.py → create_task.py}

@@ -5,9 +5,9 @@ def create_file(file_path,file_content):
     with open(file_path,"w") as f:
         f.write(file_content)
             
-def create_data(data_dir,dataset_name,sub_task_dir):
-    # call src/sub_task/scripts/create_dataset.py dataset_name
-    os.system(f"python {sub_task_dir}/scripts/create_dataset.py {dataset_name}")
+def create_data(data_dir,dataset_name,task_dir):
+    # call src/task/scripts/create_dataset.py dataset_name
+    os.system(f"python {task_dir}/scripts/create_dataset.py {dataset_name}")
 
     register_datasets_file_path = os.path.join(data_dir,"register_datasets.py")
     create_file(register_datasets_file_path,
@@ -19,7 +19,7 @@ datasets = ["{dataset_name}"]
 
 
 
-def create_model(model_dir:str, model_name:str, sub_task_dir:str):
+def create_model(model_dir:str, model_name:str, task_dir:str):
     base_model_interface_path = os.path.join(model_dir,"base_model_interface.py")
     
     create_file(base_model_interface_path,
@@ -27,26 +27,34 @@ def create_model(model_dir:str, model_name:str, sub_task_dir:str):
 from abc import ABC, abstractmethod
 
 # BaseModel Abstract class
-# all the models within this sub_task must inherit this class
+# all the models within this task must inherit this class
 
 class BaseModel(ABC):
     @abstractmethod
     def train(self):
         pass
         
+    @abstractmethod
+    def evaluate(self):
+        pass
+
     @abstractmethod        
     def predict(self,inputs):
         pass
+        
+    @abstractmethod
+    def show_results(self):
+        pass
 """)
 
 
-    # call src/sub_task/scripts/create_model.py model_name
-    os.system(f"python {sub_task_dir}/scripts/create_model.py {model_name}")
+    # call src/task/scripts/create_model.py model_name
+    os.system(f"python {task_dir}/scripts/create_model.py {model_name}")
     
 
     register_models_path = os.path.join(model_dir,"register_models.py")
     create_file(register_models_path,
-f"""# register models of this sub_task here
+f"""# register models of this task here
 models = ["{model_name}"]
 """)
 
@@ -73,7 +81,7 @@ models = ["{model_name}"]
 """)
     
 
-def create_scripts(scripts_dir,sub_task):
+def create_scripts(scripts_dir,task):
     create_dataset_path = os.path.join(scripts_dir,"create_dataset.py")
     create_file(create_dataset_path,
 f"""import os,shutil
@@ -86,7 +94,7 @@ def create_file(file_path,file_content):
 def create_dataset(args):
     dataset_name = args.name
     force_flag = args.force
-    datasets_dir = os.path.join('src','{sub_task}','data','datasets')
+    datasets_dir = os.path.join('src','{task}','data','datasets')
     
     os.makedirs(datasets_dir,exist_ok=True)
     dataset_path = os.path.join(datasets_dir,dataset_name+".py")
@@ -157,7 +165,7 @@ def create_file(file_path,file_content):
 def create_model(args):
     model_name = args.name
     force_flag = args.force
-    models_dir = os.path.join('src','{sub_task}','model',"models")
+    models_dir = os.path.join('src','{task}','model',"models")
     os.makedirs(models_dir,exist_ok=True)
     model_path = os.path.join(models_dir,model_name+".py")
 
@@ -173,7 +181,7 @@ def create_model(args):
     
     model_name_camel_case = "".join([part.capitalize() for part in model_name.split("_")])
     create_file(model_path,
-f\"\"\"from src.{sub_task}.model.base_model_interface import BaseModel
+f\"\"\"from src.{task}.model.base_model_interface import BaseModel
 
 class Model(BaseModel):
     def train(self):
@@ -197,35 +205,35 @@ if __name__=="__main__":
     
     
 
-def create_sub_task(args):
-    """Used to create sub_task within our main task"""
-    sub_task = args.sub_task
+def create_task(args):
+    """Used to create task within our main task"""
+    task = args.task
     force_flag = args.force
     dataset_name = "dataset1"
     model_name = "model1"
 
-    sub_task_dir = os.path.join('src',sub_task)
-    data_dir = os.path.join(sub_task_dir,'data')
-    model_dir = os.path.join(sub_task_dir,'model')
-    scripts_dir = os.path.join(sub_task_dir,"scripts")
+    task_dir = os.path.join('src',task)
+    data_dir = os.path.join(task_dir,'data')
+    model_dir = os.path.join(task_dir,'model')
+    scripts_dir = os.path.join(task_dir,"scripts")
     # print(scripts_dir)
-    # deleted old sub_task if force flag exists and sub_task already exists
-    if os.path.exists(sub_task_dir):
+    # deleted old task if force flag exists and task already exists
+    if os.path.exists(task_dir):
         if force_flag:
-            print("Replacing existing sub_task:",sub_task)
-            shutil.rmtree(sub_task_dir)
+            print("Replacing existing task:",task)
+            shutil.rmtree(task_dir)
         else:
-            print(f"{sub_task} already exists, use --force flag if you want to reset it to default")
+            print(f"{task} already exists, use --force flag if you want to reset it to default")
             exit()
 
     # create empty folders
-    os.makedirs(sub_task_dir,exist_ok=True)
+    os.makedirs(task_dir,exist_ok=True)
     os.makedirs(data_dir,exist_ok=True)
     os.makedirs(model_dir,exist_ok=True)
     os.makedirs(scripts_dir,exist_ok=True)
 
     # make config validator file
-    validate_config_file_path = os.path.join(sub_task_dir,"validate_config.py")
+    validate_config_file_path = os.path.join(task_dir,"validate_config.py")
     create_file(validate_config_file_path,
 '''# from cerberus import Validator
 
@@ -252,22 +260,22 @@ schema = {
 ''')
     
     # make scripts files
-    create_scripts(scripts_dir,sub_task)
+    create_scripts(scripts_dir,task)
 
     # make data files
-    create_data(data_dir,dataset_name,sub_task_dir)
+    create_data(data_dir,dataset_name,task_dir)
 
     # make model files
-    create_model(model_dir,model_name,sub_task_dir)
+    create_model(model_dir,model_name,task_dir)
 
     
 def main():
-    parser = argparse.ArgumentParser(description="Create blueprint sub_task")
-    parser.add_argument('sub_task',type=str,help="sub_task of project (e.g., simple_regression_colorization)")
-    parser.add_argument("--force",action="store_true",help="forcefully replace old existing sub_task to default",default=False)
+    parser = argparse.ArgumentParser(description="Create blueprint task")
+    parser.add_argument('task',type=str,help="task of project (e.g., simple_regression_colorization)")
+    parser.add_argument("--force",action="store_true",help="forcefully replace old existing task to default",default=False)
     args = parser.parse_args()
 
-    create_sub_task(args)
+    create_task(args)
 
 if __name__=="__main__":
     main()

src/scripts/train.py

@@ -0,0 +1,41 @@
+import os
+import argparse
+from src.utils.config_loader import Config
+from src.utils import config_loader
+from src.utils.script_utils import validate_config
+import importlib
+from pathlib import Path
+
+
+def train(args):
+    config_file_path = args.config_file
+    config = Config(config_file_path)
+
+    # validate config
+    validate_config(config)
+
+    # set config globally
+    config_loader.config = config
+
+    # now visualize the dataset
+    Model = importlib.import_module(f"src.{config.task}.model.models.{config.model}").Model
+
+
+    model_dir = os.path.join("models",config.task,config.model)
+    os.makedirs(model_dir,exist_ok=True)
+    model_save_path = os.path.join(model_dir,"model.weights.h5")
+    
+    model = Model()
+    model.train()
+    model.save(model_save_path)
+    metrics = model.evaluate()
+    print("Model Evaluation Metrics:",metrics)
+
+def main():
+    parser = argparse.ArgumentParser(description="train model based on config yaml file")
+    parser.add_argument("config_file",type=str)
+    args = parser.parse_args()
+    train(args)
+
+if __name__=="__main__":
+    main()

src/scripts/visualize_results.py

@@ -0,0 +1,41 @@
+import os
+import argparse
+from src.utils.config_loader import Config
+from src.utils import config_loader
+from src.utils.script_utils import validate_config
+import importlib
+
+
+def visualize_dataset(args):
+    config_file_path = args.config_file
+    config = Config(config_file_path)
+
+    # validate config
+    validate_config(config)
+
+    # set config globally
+    config_loader.config = config
+
+    # now load model and visualize the results
+    model_dir = os.path.join("models",config.task,config.model)
+    model_save_path = os.path.join(model_dir,"model.weights.h5")
+
+    if not os.path.exists(model_save_path):
+        raise Exception("No model found:","first use train.py to train and export a model")
+    
+    Model = importlib.import_module(f"src.{config.task}.model.models.{config.model}").Model
+    model = Model(model_save_path)
+
+    # model.train_ds
+    model.show_results()
+    
+    
+
+def main():
+    parser = argparse.ArgumentParser(description="Prepare dataset based on config yaml file")
+    parser.add_argument("config_file",type=str)
+    args = parser.parse_args()
+    visualize_dataset(args)
+
+if __name__=="__main__":
+    main()

src/simple_regression_colorization/model/base_model_interface.py

@@ -9,6 +9,14 @@ class BaseModel(ABC):
     def train(self):
         pass
         
+    @abstractmethod
+    def evaluate(self):
+        pass
+
     @abstractmethod        
     def predict(self,inputs):
         pass
+
+    @abstractmethod
+    def show_results(self):
+        pass

src/simple_regression_colorization/model/dataloaders.py

@@ -1,14 +1,14 @@
 import tensorflow as tf
 from src.utils.data_utils import scale_L,scale_AB,rescale_AB,rescale_L
-from src.utils.config_loader import config
+from src.utils.config_loader import config,constants
 from pathlib import Path
 from glob import glob
 import sklearn.model_selection
 from skimage.color import rgb2lab, lab2rgb
 
 def get_datasets():
-    trainval_dir = config.PROCESSED_DATASET_DIR / Path("trainval/")
-    test_dir = config.PROCESSED_DATASET_DIR / Path("test/")
+    trainval_dir = constants.PROCESSED_DATASET_DIR / Path("trainval/")
+    test_dir = constants.PROCESSED_DATASET_DIR / Path("test/")
 
     trainval_paths = glob(str(trainval_dir/Path("*")))
     test_paths = glob(str(test_dir/Path("*")))

src/simple_regression_colorization/model/models/model_v1.py

@@ -1,30 +1,119 @@
 from src.simple_regression_colorization.model.base_model_interface import BaseModel
 from src.simple_regression_colorization.model.dataloaders import get_datasets
+from src.utils.config_loader import config
+from src.utils.data_utils import scale_L,scale_AB,rescale_AB,rescale_L,see_batch
+from skimage.color import lab2rgb
+import tensorflow as tf
+from tensorflow.keras import layers,Model as keras_Model,Sequential
+import numpy as np
+
+def down(filters,kernel_size,apply_batch_normalization=True):
+    down = Sequential()
+    down.add(layers.Conv2D(filters,kernel_size,padding="same",strides=2))
+    if apply_batch_normalization:
+        down.add(layers.BatchNormalization())
+    down.add(layers.LeakyReLU())
+    return down
+
+def up(filters,kernel_size,dropout=False):
+    upsample = Sequential()
+    upsample.add(layers.Conv2DTranspose(filters,kernel_size,padding="same",strides=2))
+    if dropout:
+        upsample.add(layers.Dropout(dropout))
+    upsample.add(layers.LeakyReLU())
+    return upsample
 
 class Model(BaseModel):
     
-    def __init__(self):
+    def __init__(self,path=None):
         # make model architecture
         # load weights (optional)
         # create dataset loaders
         # train
         # predict
         self.init_model()
-        self.load_weights()
-        self.prepare_data()
+        self.load_weights(path)
         
 
     def init_model(self):
-        pass
+        x = layers.Input([config.image_size,config.image_size,1])
+        d1 = down(128,(3,3),False)(x)
+        d2 = down(128,(3,3),False)(d1)
+        d3 = down(256,(3,3),True)(d2)
+        d4 = down(512,(3,3),True)(d3)
+        d5 = down(512,(3,3),True)(d4)
+        
+        u1 = up(512,(3,3))(d5)
+        u1 = layers.concatenate([u1,d4])
+        u2 = up(256,(3,3))(u1)
+        u2 = layers.concatenate([u2,d3])
+        u3 = up(128,(3,3))(u2)
+        u3 = layers.concatenate([u3,d2])
+        u4 = up(128,(3,3))(u3)
+        u4 = layers.concatenate([u4,d1])
+        u5 = up(64,(3,3))(u4)
+        u5 = layers.concatenate([u5,x])
+        
+        y = layers.Conv2D(2,(2,2),strides = 1, padding = 'same',activation="tanh")(u5)
+
+        self.model = keras_Model(x,y,name="UNet")
+
 
     def load_weights(self,path=None):
-        pass
+        if path:
+            self.model.load_weights(path)
 
     def prepare_data(self):
         self.train_ds,self.val_ds,self.test_ds = get_datasets()
         
     def train(self):
-        pass
+
+        self.prepare_data()
+        self.model.compile(optimizer="adam",loss="mse",metrics=["mae","acc"])
+        self.history = self.model.fit(self.train_ds,
+                                      validation_data=self.val_ds,
+                                      epochs=config.epochs)
+
+    def save(self,model_path):
+        self.model.save_weights(model_path)
+
+    def predict(self,L_batch):
+        L_batch = scale_L(L_batch)
+        AB_batch = self.model.predict(L_batch,verbose=0)
+        return rescale_AB(AB_batch)
+
+    def evaluate(self):
+        train_metrics = self.model.evaluate(self.train_ds)
+        val_metrics = self.model.evaluate(self.val_ds)
+        test_metrics = self.model.evaluate(self.test_ds)
+
+        return {
+            "train": train_metrics,
+            "val": val_metrics,
+            "test": test_metrics,
+        }
+
+    def predict_colors(self,L_batch):
+        AB_batch = self.predict(L_batch)
+        colored_batch = np.concatenate([L_batch,rescale_AB(AB_batch)],axis=-1)
+        colored_batch = lab2rgb(colored_batch) * 255
+        return colored_batch
+
+    def show_results(self):
+        self.prepare_data()
+        
+        L_batch,AB_batch = next(iter(self.train_ds))
+        L_batch = L_batch.numpy()
+        AB_pred = self.model.predict(L_batch,verbose=0)
+        see_batch(L_batch,AB_pred,title="Train dataset Results")
+        
+        L_batch,AB_batch = next(iter(self.val_ds))
+        L_batch = L_batch.numpy()
+        AB_pred = self.model.predict(L_batch,verbose=0)
+        see_batch(L_batch,AB_pred,title="Val dataset Results")
+        
+        L_batch,AB_batch = next(iter(self.test_ds))
+        L_batch = L_batch.numpy()
+        AB_pred = self.model.predict(L_batch,verbose=0)
+        see_batch(L_batch,AB_pred,title="Test dataset Results")
         
-    def predict(self,inputs):
-        pass

src/utils/data_utils.py

@@ -75,3 +75,26 @@ def show_images_from_paths(image_paths:list[str],image_size=64,cols=4,row_size=5
             img = np.concatenate([BW,img],axis=1)
         plt.imshow(img.astype("uint8"))
     plt.show()
+
+
+def see_batch(L_batch,AB_batch,show_L=False,cols=4,row_size=5,col_size=5,title=None):
+    n = L_batch.shape[0]
+    rows = math.ceil(n/cols)
+    fig = plt.figure(figsize=(col_size*cols,row_size*rows))
+    if title:
+        plt.title(title)
+    plt.axis("off")
+    
+    for i in range(n):
+        fig.add_subplot(rows,cols,i+1)
+        L,AB = L_batch[i],AB_batch[i]
+        L,AB = rescale_L(L), rescale_AB(AB)
+#         print(L.shape,AB.shape)
+        img = np.concatenate([L,AB],axis=-1)
+        img = cv2.cvtColor(img,cv2.COLOR_LAB2RGB)*255
+#         print(img.min(),img.max())
+        if show_L:
+            L = np.tile(L,(1,1,3))/100*255
+            img = np.concatenate([L,img],axis=1)
+        plt.imshow(img.astype("uint8"))
+    plt.show()