added comet logging

src/scripts/train.py

@@ -1,11 +1,13 @@
 import os
 import argparse
+from comet_ml import Experiment
 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
-
+from dotenv import load_dotenv
+load_dotenv()
 
 def train(args):
     config_file_path = args.config_file
@@ -17,7 +19,7 @@ def train(args):
     # set config globally
     config_loader.config = config
 
-    # now visualize the dataset
+    # now load the model
     Model = importlib.import_module(f"src.{config.task}.model.models.{config.model}").Model
 
 
@@ -25,11 +27,30 @@ def train(args):
     os.makedirs(model_dir,exist_ok=True)
     model_save_path = os.path.join(model_dir,"model.weights.h5")
     
-    model = Model()
+
+    experiment = Experiment(
+        api_key=os.environ["COMET_API_KEY"],
+        project_name="image-colorization",
+        workspace="anujpanthri",
+        auto_histogram_activation_logging=True,
+        auto_histogram_epoch_rate=True,
+        auto_histogram_gradient_logging=True,
+        auto_histogram_weight_logging=True,
+        auto_param_logging=True,
+    )
+
+    model = Model(experiment=experiment)
     model.train()
     model.save(model_save_path)
+    
+    # log model to comet
+    if "LOCAL_SYSTEM" not in os.environ:
+        experiment.log_model(f"{config.task}_{config.dataset}_{config.model}",model_save_path)
+
     metrics = model.evaluate()
     print("Model Evaluation Metrics:",metrics)
+    
+    experiment.end()
 
 def main():
     parser = argparse.ArgumentParser(description="train model based on config yaml file")

src/simple_regression_colorization/model/base_model_interface.py

@@ -1,22 +1,88 @@
 import numpy as np
 from abc import ABC, abstractmethod
+from src.simple_regression_colorization.model.dataloaders import get_datasets
+from src.utils.data_utils import scale_L,scale_AB,rescale_AB,rescale_L,see_batch
+from src.utils.config_loader import config
+from skimage.color import lab2rgb
+from src.simple_regression_colorization.model.callbacks import LogPredictionsCallback
 
 # BaseModel Abstract class
 # all the models within this sub_task must inherit this class
 
 class BaseModel(ABC):
-    @abstractmethod
+    def __init__(self,path=None,experiment=None):
+        self.init_model()
+        self.load_weights(path)
+        self.experiment = experiment
+    
+    def load_weights(self,path=None):
+        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"])
         
-    @abstractmethod
+        callbacks = [
+            LogPredictionsCallback(self.train_ds,"train_ds",self.experiment),
+            LogPredictionsCallback(self.val_ds,"val_ds",self.experiment),
+        ]
+
+        self.history = self.model.fit(self.train_ds,
+                                      validation_data=self.val_ds,
+                                      callbacks=callbacks,
+                                      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):
-        pass
+        train_metrics = self.model.evaluate(self.train_ds)
+        val_metrics = self.model.evaluate(self.val_ds)
+        test_metrics = self.model.evaluate(self.test_ds)
 
-    @abstractmethod        
-    def predict(self,inputs):
-        pass
+        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
 
-    @abstractmethod
     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")
+        
+
+    @abstractmethod
+    def init_model(self):
         pass
+        
+  

src/simple_regression_colorization/model/callbacks.py

@@ -1 +1,27 @@
-# define callbacks here
+import tensorflow as tf
+from src.utils.data_utils import rescale_AB,rescale_L
+from skimage.color import lab2rgb
+import numpy as np
+
+class LogPredictionsCallback(tf.keras.callbacks.Callback):
+    def __init__(self,ds,ds_name,experiment=None):
+        self.ds = ds
+        self.ds_name = ds_name
+        self.experiment = experiment
+
+    def on_epoch_end(self, epoch, logs=None):
+        
+        L_batch, _ = next(iter(self.ds))
+        AB_batch = self.model.predict(L_batch,verbose=0)
+        colored_batch = np.concatenate([rescale_L(L_batch),rescale_AB(AB_batch)],axis=-1)
+        colored_batch = lab2rgb(colored_batch) * 255
+        
+        print(self.ds_name)
+        print("R:",colored_batch[:,:,0].min(),colored_batch[:,:,0].max())
+        print("G:",colored_batch[:,:,1].min(),colored_batch[:,:,1].max())
+        print("B:",colored_batch[:,:,2].min(),colored_batch[:,:,2].max())
+        
+        if self.experiment:
+            # log images
+            for i,image in enumerate(colored_batch):
+                self.experiment.log_image(image,name=f"{self.ds_name}_{i}")

src/simple_regression_colorization/model/model_utils.py

@@ -0,0 +1,18 @@
+import tensorflow as tf
+from tensorflow.keras import layers,Model as keras_Model,Sequential
+
+def down_block(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_block(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

src/simple_regression_colorization/model/models/model_v1.py

@@ -1,119 +1,34 @@
 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
+from src.simple_regression_colorization.model.model_utils import up_block,down_block
 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,path=None):
-        # make model architecture
-        # load weights (optional)
-        # create dataset loaders
-        # train
-        # predict
-        self.init_model()
-        self.load_weights(path)
-        
+    def __init__(self,*args,**kwargs):
+        super().__init__(*args,**kwargs)
 
     def init_model(self):
         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)
+        d1 = down_block(128,(3,3),False)(x)
+        d2 = down_block(128,(3,3),False)(d1)
+        d3 = down_block(256,(3,3),True)(d2)
+        d4 = down_block(512,(3,3),True)(d3)
+        d5 = down_block(512,(3,3),True)(d4)
         
-        u1 = up(512,(3,3))(d5)
+        u1 = up_block(512,(3,3))(d5)
         u1 = layers.concatenate([u1,d4])
-        u2 = up(256,(3,3))(u1)
+        u2 = up_block(256,(3,3))(u1)
         u2 = layers.concatenate([u2,d3])
-        u3 = up(128,(3,3))(u2)
+        u3 = up_block(128,(3,3))(u2)
         u3 = layers.concatenate([u3,d2])
-        u4 = up(128,(3,3))(u3)
+        u4 = up_block(128,(3,3))(u3)
         u4 = layers.concatenate([u4,d1])
-        u5 = up(64,(3,3))(u4)
+        u5 = up_block(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):
-        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):
-
-        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")
-        
+        self.model = keras_Model(x,y,name="UNet")