made some improvement

.gitignore

@@ -88,4 +88,5 @@ target/
 # Mypy cache
 .mypy_cache/
 
-/models
+/models
+/artifacts

configs/experiment1.yaml

@@ -12,5 +12,5 @@ shuffle: False
 
 # training related
 batch_size: 16
-epochs: 15
-# epochs: 02
+# epochs: 15
+epochs: 02

constants.yaml

@@ -1,3 +1,8 @@
 RAW_DATASET_DIR: data/raw/
 INTERIM_DATASET_DIR: data/interim/
-PROCESSED_DATASET_DIR: data/processed/
+PROCESSED_DATASET_DIR: data/processed/
+
+
+ARTIFACT_MODEL_DIR: artifacts/model/
+ARTIFACT_DATASET_VISUALIZATION_DIR: artifacts/dataset/
+ARTIFACT_RESULT_VISUALIZATION_DIR: artifacts/result/

src/scripts/train.py

@@ -1,7 +1,7 @@
 import os,shutil
 import argparse
 from comet_ml import Experiment
-from src.utils.config_loader import Config
+from src.utils.config_loader import Config,constants
 from src.utils import config_loader
 from src.utils.data_utils import print_title
 from src.utils.script_utils import validate_config
@@ -24,42 +24,51 @@ def train(args):
     Model = importlib.import_module(f"src.{config.task}.model.models.{config.model}").Model
 
 
-    model_dir = os.path.join("models",config.task,config.model)
+    model_dir = constants.ARTIFACT_MODEL_DIR
     os.makedirs(model_dir,exist_ok=True)
     model_save_path = os.path.join(model_dir,"model.weights.h5")
     
     # save config to exported model folder
     shutil.copy(config_file_path,model_dir)
+    # rename it to config.yaml
+    shutil.move(os.path.join(model_dir,Path(config_file_path).name),os.path.join(model_dir,"config.yaml"))
 
-    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,
-    )
+    experiment = None
+    if args.log:
+        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)
+
+    print_title("\nTraining Model")
     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_dir)
-
+        if experiment:
+            experiment.log_model(f"{config.task}_{config.dataset}_{config.model}",model_dir)
+        
     # evaluate model
     print_title("\nEvaluating Model")
     metrics = model.evaluate()
     print("Model Evaluation Metrics:",metrics)
     
-    experiment.end()
+    if experiment:
+        experiment.end()
 
 def main():
     parser = argparse.ArgumentParser(description="train model based on config yaml file")
     parser.add_argument("config_file",type=str)
+    parser.add_argument("--log",action="store_true",default=False)
     args = parser.parse_args()
     train(args)
 

src/scripts/visualize_results.py

@@ -1,12 +1,12 @@
 import os
 import argparse
-from src.utils.config_loader import Config
+from src.utils.config_loader import Config,constants
 from src.utils import config_loader
 from src.utils.script_utils import validate_config
 import importlib
 
 
-def visualize_dataset(args):
+def visualize_results(args):
     config_file_path = args.config_file
     config = Config(config_file_path)
 
@@ -17,7 +17,7 @@ def visualize_dataset(args):
     config_loader.config = config
 
     # now load model and visualize the results
-    model_dir = os.path.join("models",config.task,config.model)
+    model_dir = constants.ARTIFACT_MODEL_DIR
     model_save_path = os.path.join(model_dir,"model.weights.h5")
 
     if not os.path.exists(model_save_path):
@@ -32,10 +32,10 @@ def visualize_dataset(args):
     
 
 def main():
-    parser = argparse.ArgumentParser(description="Prepare dataset based on config yaml file")
+    parser = argparse.ArgumentParser(description="visualize results based on config yaml file and trained model")
     parser.add_argument("config_file",type=str)
     args = parser.parse_args()
-    visualize_dataset(args)
+    visualize_results(args)
 
 if __name__=="__main__":
     main()

src/simple_regression_colorization/data/visualize_dataset.py

@@ -12,10 +12,16 @@ def visualize():
     choosen_paths = np.random.choice(image_paths,n)
     show_images_from_paths(choosen_paths,
                            title="sample of train_val dataset",
-                           image_size=config.image_size)
+                           image_size=config.image_size,
+                           save=True,
+                           label="trainval",
+                           )
 
     image_paths = glob(f"{constants.PROCESSED_DATASET_DIR}/test/*")
     choosen_paths = np.random.choice(image_paths,n)
     show_images_from_paths(choosen_paths,
                            title="sample of test dataset",
-                           image_size=config.image_size)
+                           image_size=config.image_size,
+                           save=True,
+                           label="test",
+                           )

src/simple_regression_colorization/model/base_model_interface.py

@@ -68,17 +68,32 @@ class BaseModel(ABC):
         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")
+        see_batch(L_batch,
+                  AB_pred,
+                  title="Train dataset Results",
+                  save = True,
+                  label = "train",
+                  )
         
         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")
+        see_batch(L_batch,
+                  AB_pred,
+                  title="Val dataset Results",
+                  save = True,
+                  label = "val",
+                  )
         
         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")
+        see_batch(L_batch,
+                  AB_pred,
+                  title="Test dataset Results",
+                  save = True,
+                  label = "test",
+                  )
         
 
     @abstractmethod

src/utils/data_utils.py

@@ -2,101 +2,148 @@ from src.utils.config_loader import constants
 from huggingface_hub import snapshot_download
 from zipfile import ZipFile
 import numpy as np
-import os,shutil
+import os, shutil
 import matplotlib.pyplot as plt
 import cv2
 import math
 
 
-def download_hf_dataset(repo_id,allow_patterns=None):
+def download_hf_dataset(repo_id, allow_patterns=None):
     """Used to download dataset from any public hugging face dataset"""
-    snapshot_download(repo_id=repo_id,
-                    repo_type="dataset",
-                    local_dir=constants.RAW_DATASET_DIR,
-                    allow_patterns=allow_patterns)
+    snapshot_download(
+        repo_id=repo_id,
+        repo_type="dataset",
+        local_dir=constants.RAW_DATASET_DIR,
+        allow_patterns=allow_patterns,
+    )
 
 
 def download_personal_hf_dataset(name):
     """Used to download dataset from a specific hugging face dataset"""
-    download_hf_dataset(repo_id="Anuj-Panthri/Image-Colorization-Datasets",
-                        allow_patterns=f"{name}/*")
+    download_hf_dataset(
+        repo_id="Anuj-Panthri/Image-Colorization-Datasets", allow_patterns=f"{name}/*"
+    )
 
 
-def unzip_file(file_path,destination_dir):
+def unzip_file(file_path, destination_dir):
     """unzips file to destination_dir"""
     if os.path.exists(destination_dir):
         shutil.rmtree(destination_dir)
     os.makedirs(destination_dir)
-    with ZipFile(file_path,"r") as zip:
+    with ZipFile(file_path, "r") as zip:
         zip.extractall(destination_dir)
 
-def is_bw(img:np.ndarray):
+
+def is_bw(img: np.ndarray):
     """checks if RGB image is black and white"""
-    rg,gb,rb = img[:,:,0]-img[:,:,1] , img[:,:,1]-img[:,:,2] , img[:,:,0]-img[:,:,2]
-    rg,gb,rb = np.abs(rg).sum(),np.abs(gb).sum(),np.abs(rb).sum()
-    avg = np.mean([rg,gb,rb])
-    
-    return avg<10
+    rg, gb, rb = (
+        img[:, :, 0] - img[:, :, 1],
+        img[:, :, 1] - img[:, :, 2],
+        img[:, :, 0] - img[:, :, 2],
+    )
+    rg, gb, rb = np.abs(rg).sum(), np.abs(gb).sum(), np.abs(rb).sum()
+    avg = np.mean([rg, gb, rb])
+
+    return avg < 10
+
 
+def print_title(msg: str, max_chars=105):
+    n = (max_chars - len(msg)) // 2
+    print("=" * n, msg.upper(), "=" * n, sep="")
 
-def print_title(msg:str,max_chars=105):
-    n = (max_chars-len(msg))//2
-    print("="*n,msg.upper(),"="*n,sep="")
 
 def scale_L(L):
-    return L/100
+    return L / 100
+
 
 def rescale_L(L):
-    return L*100
+    return L * 100
+
 
 def scale_AB(AB):
-    return AB/128
+    return AB / 128
+
 
 def rescale_AB(AB):
-    return AB*128
-    
+    return AB * 128
+
 
+def show_images_from_paths(
+    image_paths: list[str],
+    image_size=64,
+    cols=4,
+    row_size=5,
+    col_size=5,
+    show_BW=False,
+    title=None,
+    save=False,
+    label="",
+):
 
-def show_images_from_paths(image_paths:list[str],image_size=64,cols=4,row_size=5,col_size=5,show_BW=False,title=None):
     n = len(image_paths)
-    rows = math.ceil(n/cols)
-    fig = plt.figure(figsize=(col_size*cols,row_size*rows))
+    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)
-        
-        img = cv2.imread(image_paths[i])[:,:,::-1]
-        img = cv2.resize(img,[image_size,image_size])
+        fig.add_subplot(rows, cols, i + 1)
+
+        img = cv2.imread(image_paths[i])[:, :, ::-1]
+        img = cv2.resize(img, [image_size, image_size])
 
         if show_BW:
-            BW = cv2.cvtColor(img,cv2.COLOR_RGB2GRAY)
-            BW = np.tile(BW,(1,1,3))
-            img = np.concatenate([BW,img],axis=1)
+            BW = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+            BW = np.tile(BW, (1, 1, 3))
+            img = np.concatenate([BW, img], axis=1)
         plt.imshow(img.astype("uint8"))
+
+    if save:
+        os.makedirs(constants.ARTIFACT_DATASET_VISUALIZATION_DIR, exist_ok=True)
+        plt.savefig(
+            os.path.join(
+                constants.ARTIFACT_DATASET_VISUALIZATION_DIR, f"{label}_image.png"
+            )
+        )
     plt.show()
 
 
-def see_batch(L_batch,AB_batch,show_L=False,cols=4,row_size=5,col_size=5,title=None):
+def see_batch(
+    L_batch,
+    AB_batch,
+    show_L=False,
+    cols=4,
+    row_size=5,
+    col_size=5,
+    title=None,
+    save=False,
+    label="",
+):
     n = L_batch.shape[0]
-    rows = math.ceil(n/cols)
-    fig = plt.figure(figsize=(col_size*cols,row_size*rows))
+    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())
+        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)
+            L = np.tile(L, (1, 1, 3)) / 100 * 255
+            img = np.concatenate([L, img], axis=1)
         plt.imshow(img.astype("uint8"))
+    if save:
+        os.makedirs(constants.ARTIFACT_RESULT_VISUALIZATION_DIR, exist_ok=True)
+        plt.savefig(
+            os.path.join(
+                constants.ARTIFACT_RESULT_VISUALIZATION_DIR, f"{label}_image.png"
+            )
+        )
     plt.show()