model building added

config/config.yaml

@@ -12,4 +12,7 @@ data_transformation:
   data_path_black : artifacts/data_ingestion/ab/ab/ab1.npy
   data_path_grey : artifacts/data_ingestion/l/gray_scale.npy
 
+model_building:
+  root_dir : artifacts/model
+  
   

main.py

@@ -1,5 +1,6 @@
 from src.imagecolorization.pipeline.stage01_data_ingestion import DataIngestionPipeline
 from src.imagecolorization.pipeline.stage02_data_transformation import DataTransformationPipeline
+from src.imagecolorization.pipeline.stage_03_model_building import ModelBuildingPipeline
 from src.imagecolorization.logging import logger
 
 STAGE_NAME = 'Data Ingestion Config'
@@ -20,6 +21,18 @@ try:
    data_transformation = DataTransformationPipeline()
    data_transformation.main()
    logger.info(f">>>>>> stage {STAGE_NAME} completed <<<<<<\n\nx==========x")
+except Exception as e:
+        logger.exception(e)
+        raise e
+     
+     
+STAGE_NAME = 'Model Building Config'
+
+try:
+   logger.info(f">>>>>> stage {STAGE_NAME} started <<<<<<") 
+   model_building = ModelBuildingPipeline()
+   model_building.main()
+   logger.info(f">>>>>> stage {STAGE_NAME} completed <<<<<<\n\nx==========x")
 except Exception as e:
         logger.exception(e)
         raise e

params.yaml

@@ -1,5 +1,25 @@
-BATCH_SIZE : 1
-IMAGE_SIZE : [224,224,1]
+# Data Parameters
+BATCH_SIZE: 1
+IMAGE_SIZE: [224, 224, 1]
 DATA_RANGE: 5000
-KERNEL_SIZE : 3
-p
+
+# Convolutional Layer Parameters
+KERNEL_SIZE_RES: 3
+PADDING: 1
+STRIDE: 1
+BIAS: False
+
+# UpSampling Parameters
+SCALE_FACTOR: 2
+DIM: 1
+
+# Dropout Parameters
+DROPOUT_RATE: 0.2
+
+# Generator Parameters
+KERNEL_SIZE_GENERATOR: 1
+INPUT_CHANNELS: 1
+OUTPUT_CHANNELS: 2
+
+# Critic Parameters
+IN_CHANNELS: 3

research/data_transformation.ipynb

@@ -164,15 +164,28 @@
     "        return train_loader, test_loader\n",
     "    \n",
     "    \n",
+    "    \n",
+    "        \n",
+    "        \n",
+    "    \n",
+    "    \n",
     "    def save_dataloaders(self, train_loader, test_loader):\n",
+    "        # Ensure the directory exists\n",
+    "        os.makedirs(self.config.root_dir, exist_ok=True)\n",
+    "\n",
     "        train_loader_path = os.path.join(self.config.root_dir, 'train_loader.pt')\n",
     "        test_loader_path = os.path.join(self.config.root_dir, 'test_loader.pt')\n",
-    "    \n",
-    "        torch.save(train_loader, train_loader_path)\n",
-    "        torch.save(test_loader, test_loader_path)\n",
-    "    \n",
-    "        logger.info(f\"Train Loader saved at: {train_loader_path}\")\n",
-    "        logger.info(f\"Test Loader saved at: {test_loader_path}\")"
+    "\n",
+    "        try:\n",
+    "            # Save the dataloaders\n",
+    "            torch.save(train_loader, train_loader_path)\n",
+    "            torch.save(test_loader, test_loader_path)\n",
+    "\n",
+    "            logger.info(f\"Train Loader saved at: {train_loader_path}\")\n",
+    "            logger.info(f\"Test Loader saved at: {test_loader_path}\")\n",
+    "        except Exception as e:\n",
+    "            logger.error(f\"Error saving dataloaders: {str(e)}\")\n",
+    "            raise e\n"
    ]
   },
   {
@@ -184,55 +197,11 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "[2024-08-18 13:20:34,232: INFO: common: yaml file: config\\config.yaml loaded successfully]\n",
-      "[2024-08-18 13:20:34,234: INFO: common: yaml file: params.yaml loaded successfully]\n",
-      "[2024-08-18 13:20:34,235: INFO: common: created directory at: artifacts]\n"
-     ]
-    },
-    {
-     "ename": "BoxKeyError",
-     "evalue": "\"'ConfigBox' object has no attribute 'data_transformation'\"",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[1;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
-      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\box\\box.py:503\u001b[0m, in \u001b[0;36mBox.__getitem__\u001b[1;34m(self, item, _ignore_default)\u001b[0m\n\u001b[0;32m    502\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m--> 503\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__getitem__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mitem\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    504\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m err:\n",
-      "\u001b[1;31mKeyError\u001b[0m: 'data_transformation'",
-      "\nThe above exception was the direct cause of the following exception:\n",
-      "\u001b[1;31mBoxKeyError\u001b[0m                               Traceback (most recent call last)",
-      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\box\\box.py:536\u001b[0m, in \u001b[0;36mBox.__getattr__\u001b[1;34m(self, item)\u001b[0m\n\u001b[0;32m    535\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m--> 536\u001b[0m     value \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__getitem__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mitem\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m_ignore_default\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[0;32m    537\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyError\u001b[39;00m:\n",
-      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\box\\box.py:524\u001b[0m, in \u001b[0;36mBox.__getitem__\u001b[1;34m(self, item, _ignore_default)\u001b[0m\n\u001b[0;32m    523\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m__get_default(item)\n\u001b[1;32m--> 524\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m BoxKeyError(\u001b[38;5;28mstr\u001b[39m(err)) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m_exception_cause\u001b[39;00m(err)\n\u001b[0;32m    525\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m err:\n",
-      "\u001b[1;31mBoxKeyError\u001b[0m: \"'data_transformation'\"",
-      "\nDuring handling of the above exception, another exception occurred:\n",
-      "\u001b[1;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
-      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\box\\box.py:538\u001b[0m, in \u001b[0;36mBox.__getattr__\u001b[1;34m(self, item)\u001b[0m\n\u001b[0;32m    537\u001b[0m     \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyError\u001b[39;00m:\n\u001b[1;32m--> 538\u001b[0m         value \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mobject\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__getattribute__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mitem\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    539\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mAttributeError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m err:\n",
-      "\u001b[1;31mAttributeError\u001b[0m: 'ConfigBox' object has no attribute 'data_transformation'",
-      "\nThe above exception was the direct cause of the following exception:\n",
-      "\u001b[1;31mBoxKeyError\u001b[0m                               Traceback (most recent call last)",
-      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\box\\config_box.py:28\u001b[0m, in \u001b[0;36mConfigBox.__getattr__\u001b[1;34m(self, item)\u001b[0m\n\u001b[0;32m     27\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m---> 28\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__getattr__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mitem\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     29\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mAttributeError\u001b[39;00m:\n",
-      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\box\\box.py:552\u001b[0m, in \u001b[0;36mBox.__getattr__\u001b[1;34m(self, item)\u001b[0m\n\u001b[0;32m    551\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m__get_default(item, attr\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m--> 552\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m BoxKeyError(\u001b[38;5;28mstr\u001b[39m(err)) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m_exception_cause\u001b[39;00m(err)\n\u001b[0;32m    553\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m value\n",
-      "\u001b[1;31mBoxKeyError\u001b[0m: \"'ConfigBox' object has no attribute 'data_transformation'\"",
-      "\nDuring handling of the above exception, another exception occurred:\n",
-      "\u001b[1;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
-      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\box\\box.py:503\u001b[0m, in \u001b[0;36mBox.__getitem__\u001b[1;34m(self, item, _ignore_default)\u001b[0m\n\u001b[0;32m    502\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m--> 503\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__getitem__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mitem\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    504\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m err:\n",
-      "\u001b[1;31mKeyError\u001b[0m: 'data_transformation'",
-      "\nThe above exception was the direct cause of the following exception:\n",
-      "\u001b[1;31mBoxKeyError\u001b[0m                               Traceback (most recent call last)",
-      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\box\\box.py:536\u001b[0m, in \u001b[0;36mBox.__getattr__\u001b[1;34m(self, item)\u001b[0m\n\u001b[0;32m    535\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m--> 536\u001b[0m     value \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__getitem__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mitem\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m_ignore_default\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[0;32m    537\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyError\u001b[39;00m:\n",
-      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\box\\box.py:524\u001b[0m, in \u001b[0;36mBox.__getitem__\u001b[1;34m(self, item, _ignore_default)\u001b[0m\n\u001b[0;32m    523\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m__get_default(item)\n\u001b[1;32m--> 524\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m BoxKeyError(\u001b[38;5;28mstr\u001b[39m(err)) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m_exception_cause\u001b[39;00m(err)\n\u001b[0;32m    525\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m err:\n",
-      "\u001b[1;31mBoxKeyError\u001b[0m: \"'data_transformation'\"",
-      "\nDuring handling of the above exception, another exception occurred:\n",
-      "\u001b[1;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
-      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\box\\box.py:538\u001b[0m, in \u001b[0;36mBox.__getattr__\u001b[1;34m(self, item)\u001b[0m\n\u001b[0;32m    537\u001b[0m     \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyError\u001b[39;00m:\n\u001b[1;32m--> 538\u001b[0m         value \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mobject\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__getattribute__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mitem\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    539\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mAttributeError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m err:\n",
-      "\u001b[1;31mAttributeError\u001b[0m: 'ConfigBox' object has no attribute 'data_transformation'",
-      "\nThe above exception was the direct cause of the following exception:\n",
-      "\u001b[1;31mBoxKeyError\u001b[0m                               Traceback (most recent call last)",
-      "Cell \u001b[1;32mIn[7], line 9\u001b[0m\n\u001b[0;32m      7\u001b[0m     data_transformation\u001b[38;5;241m.\u001b[39msa\n\u001b[0;32m      8\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m----> 9\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m e\n",
-      "Cell \u001b[1;32mIn[7], line 3\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m      2\u001b[0m     config \u001b[38;5;241m=\u001b[39m ConfigurationManager()\n\u001b[1;32m----> 3\u001b[0m     data_transformation_config \u001b[38;5;241m=\u001b[39m \u001b[43mconfig\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_data_transformation_config\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m      4\u001b[0m     data_transformation \u001b[38;5;241m=\u001b[39m DataTransformation(config\u001b[38;5;241m=\u001b[39mdata_transformation_config)\n\u001b[0;32m      5\u001b[0m     data_transformation\u001b[38;5;241m.\u001b[39mload_data()\n",
-      "Cell \u001b[1;32mIn[4], line 17\u001b[0m, in \u001b[0;36mConfigurationManager.get_data_transformation_config\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m     16\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mget_data_transformation_config\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m DataTransformationConfig:\n\u001b[1;32m---> 17\u001b[0m     config \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconfig\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdata_transformation\u001b[49m\n\u001b[0;32m     18\u001b[0m     params \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mparams\n\u001b[0;32m     20\u001b[0m     data_transformation_cofig \u001b[38;5;241m=\u001b[39m DataTransformationConfig(\n\u001b[0;32m     21\u001b[0m         root_dir\u001b[38;5;241m=\u001b[39m config\u001b[38;5;241m.\u001b[39mroot_dir,\n\u001b[0;32m     22\u001b[0m         data_path_black\u001b[38;5;241m=\u001b[39mconfig\u001b[38;5;241m.\u001b[39mdata_path_black,\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m     26\u001b[0m         DATA_RANGE\u001b[38;5;241m=\u001b[39mparams\u001b[38;5;241m.\u001b[39mDATA_RANGE\n\u001b[0;32m     27\u001b[0m     )\n",
-      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\box\\config_box.py:30\u001b[0m, in \u001b[0;36mConfigBox.__getattr__\u001b[1;34m(self, item)\u001b[0m\n\u001b[0;32m     28\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28msuper\u001b[39m()\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__getattr__\u001b[39m(item)\n\u001b[0;32m     29\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mAttributeError\u001b[39;00m:\n\u001b[1;32m---> 30\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;21;43m__getattr__\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mitem\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlower\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
-      "File \u001b[1;32mc:\\Users\\azizu\\AppData\\Local\\Programs\\Python\\Python311\\Lib\\site-packages\\box\\box.py:552\u001b[0m, in \u001b[0;36mBox.__getattr__\u001b[1;34m(self, item)\u001b[0m\n\u001b[0;32m    550\u001b[0m             \u001b[38;5;28;01mraise\u001b[39;00m BoxKeyError(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mitem\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m: Does not exist and internal methods are never defaulted\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m    551\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m__get_default(item, attr\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m--> 552\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m BoxKeyError(\u001b[38;5;28mstr\u001b[39m(err)) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m_exception_cause\u001b[39;00m(err)\n\u001b[0;32m    553\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m value\n",
-      "\u001b[1;31mBoxKeyError\u001b[0m: \"'ConfigBox' object has no attribute 'data_transformation'\""
+      "[2024-08-18 17:50:45,127: INFO: common: yaml file: config\\config.yaml loaded successfully]\n",
+      "[2024-08-18 17:50:45,129: INFO: common: yaml file: params.yaml loaded successfully]\n",
+      "[2024-08-18 17:50:45,129: INFO: common: created directory at: artifacts]\n",
+      "[2024-08-18 17:50:57,600: INFO: 2567581832: Train Loader saved at: artifacts/data_transformation\\train_loader.pt]\n",
+      "[2024-08-18 17:50:57,605: INFO: 2567581832: Test Loader saved at: artifacts/data_transformation\\test_loader.pt]\n"
      ]
     }
    ],
@@ -241,11 +210,18 @@
     "    config = ConfigurationManager()\n",
     "    data_transformation_config = config.get_data_transformation_config()\n",
     "    data_transformation = DataTransformation(config=data_transformation_config)\n",
-    "    data_transformation.load_data()\n",
-    "    data_transformation.get_dataloader()\n",
-    "    data_transformation.sa\n",
+    "    \n",
+    "    # Load the dataset\n",
+    "    dataset = data_transformation.load_data()\n",
+    "    \n",
+    "    # Get the dataloader using the loaded dataset\n",
+    "    train_loader, test_loader = data_transformation.get_dataloader(dataset)\n",
+    "    \n",
+    "    # Perform any further operations (e.g., saving the dataloaders)\n",
+    "    data_transformation.save_dataloaders(train_loader, test_loader)\n",
+    "    \n",
     "except Exception as e:\n",
-    "    raise e"
+    "    raise e\n"
    ]
   },
   {

research/model_building.ipynb

@@ -0,0 +1,461 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "os.chdir('../')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'c:\\\\mlops project\\\\image-colorization-mlops'"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "%pwd"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from dataclasses import dataclass\n",
+    "from pathlib import Path\n",
+    "\n",
+    "@dataclass(frozen=True)\n",
+    "class ModelBuildingConfig:\n",
+    "    root_dir: Path\n",
+    "    KERNEL_SIZE_RES: int\n",
+    "    PADDING: int\n",
+    "    STRIDE: int\n",
+    "    BIAS: bool\n",
+    "    SCALE_FACTOR: int\n",
+    "    DIM: int\n",
+    "    DROPOUT_RATE: float\n",
+    "    KERNEL_SIZE_GENERATOR: int\n",
+    "    INPUT_CHANNELS: int\n",
+    "    OUTPUT_CHANNELS: int\n",
+    "    IN_CHANNELS: int\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from src.imagecolorization.constants import *\n",
+    "from src.imagecolorization.utils.common import read_yaml, create_directories\n",
+    "\n",
+    "class ConfigurationManager:\n",
+    "    def __init__(self, config_filepath=CONFIG_FILE_PATH, params_filepath=PARAMS_FILE_PATH):\n",
+    "        self.config = read_yaml(config_filepath)\n",
+    "        self.params = read_yaml(params_filepath)\n",
+    "        create_directories([self.config.artifacts_root])\n",
+    "\n",
+    "    def get_model_building_config(self) -> ModelBuildingConfig:\n",
+    "        config = self.config.model_building\n",
+    "        params = self.params\n",
+    "\n",
+    "        model_building_config = ModelBuildingConfig(\n",
+    "            root_dir=Path(config.root_dir),\n",
+    "            KERNEL_SIZE_RES=params.KERNEL_SIZE_RES,\n",
+    "            PADDING=params.PADDING,\n",
+    "            STRIDE=params.STRIDE,\n",
+    "            BIAS=params.BIAS,\n",
+    "            SCALE_FACTOR=params.SCALE_FACTOR,\n",
+    "            DIM=params.DIM,\n",
+    "            DROPOUT_RATE=params.DROPOUT_RATE,\n",
+    "            KERNEL_SIZE_GENERATOR=params.KERNEL_SIZE_GENERATOR,\n",
+    "            INPUT_CHANNELS=params.INPUT_CHANNELS,\n",
+    "            OUTPUT_CHANNELS=params.OUTPUT_CHANNELS,\n",
+    "            IN_CHANNELS=params.IN_CHANNELS\n",
+    "        )\n",
+    "        return model_building_config\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch \n",
+    "import torch.nn as nn\n",
+    "from pathlib import Path\n",
+    "\n",
+    "class ResBlock(nn.Module):\n",
+    "    def __init__(self, in_channles, out_channels, stride = 1, kerenl_size = 3, padding = 1, bias = False):\n",
+    "        super().__init__()\n",
+    "        self.layer = nn.Sequential(\n",
+    "            nn.Conv2d(in_channles, out_channels, kernel_size=kerenl_size, padding=padding, stride=stride, bias = bias),\n",
+    "            nn.BatchNorm2d(out_channels),\n",
+    "            nn.ReLU(inplace=True),\n",
+    "            nn.Conv2d(out_channels, out_channels, kernel_size=kerenl_size, padding=padding, stride = 1, bias = bias),\n",
+    "            nn.BatchNorm2d(out_channels),\n",
+    "            nn.ReLU(inplace=True)\n",
+    "        )\n",
+    "        \n",
+    "        self.identity_map = nn.Conv2d(in_channles, out_channels,kernel_size=1, stride=stride)\n",
+    "        self.relu = nn.ReLU(inplace= True)\n",
+    "        \n",
+    "    def forward(self, inputs):\n",
+    "        x = inputs.clone().detach()\n",
+    "        out = self.layer(x)\n",
+    "        residual = self.identity_map(inputs)\n",
+    "        skip  = out + residual\n",
+    "        return self.relu(skip)\n",
+    "    \n",
+    "    \n",
+    "class DownsampleConv(nn.Module):\n",
+    "    def __init__(self, in_channels, out_channels, stride = 1):\n",
+    "        super().__init__()\n",
+    "        self.layer = nn.Sequential(\n",
+    "            nn.MaxPool2d(2),\n",
+    "            ResBlock(in_channels, out_channels)\n",
+    "        )\n",
+    "        \n",
+    "    def forward(self, inputs):\n",
+    "        return self.layer(inputs)\n",
+    "    \n",
+    "    \n",
+    "    \n",
+    "class UpsampleConv(nn.Module):\n",
+    "    def __init__(self, in_channels, out_channels, scale_factor=2):\n",
+    "        super().__init__()\n",
+    "        self.upsample = nn.Upsample(scale_factor=scale_factor,mode = 'bilinear', align_corners=True)\n",
+    "        self.res_block = ResBlock(in_channels + out_channels, out_channels)\n",
+    "\n",
+    "    def forward(self, inputs, skip):\n",
+    "        x = self.upsample(inputs)\n",
+    "        x = torch.cat([x, skip], dim = 1)\n",
+    "        x = self.res_block(x)\n",
+    "        return x\n",
+    "    \n",
+    "class Generator(nn.Module):\n",
+    "    def __init__(self, input_channels, output_channels, dropout_rate = 0.2):\n",
+    "        super().__init__()\n",
+    "        self.encoding_layer1_= ResBlock(input_channels, 64)\n",
+    "        self.encoding_layer2_ = DownsampleConv(64, 128)\n",
+    "        self.encoding_layer3_ = DownsampleConv(128, 256)\n",
+    "        self.bridge = DownsampleConv(256, 512)\n",
+    "        self.decoding_layer3 = UpsampleConv(512, 256)\n",
+    "        self.decoding_layer2 = UpsampleConv(256, 128)\n",
+    "        self.decoding_layer1 = UpsampleConv(128 , 64)\n",
+    "        self.output = nn.Conv2d(64, output_channels, kernel_size = 1)\n",
+    "        self.dropout = nn.Dropout2d(dropout_rate)\n",
+    "        \n",
+    "    def forward(self, inputs):\n",
+    "        e1 = self.encoding_layer1_(inputs)\n",
+    "        e1 = self.dropout(e1)\n",
+    "        e2 = self.encoding_layer2_(e1)\n",
+    "        e2 = self.dropout(e2)\n",
+    "        e3 = self.encoding_layer3_(e2)\n",
+    "        e3 = self.dropout(e3)\n",
+    "        \n",
+    "        bridge = self.bridge(e3)\n",
+    "        bridge = self.dropout(bridge)\n",
+    "        \n",
+    "        d3 = self.decoding_layer3(bridge, e3)\n",
+    "        d2 =self.decoding_layer2(d3, e2)\n",
+    "        d1 = self.decoding_layer1(d2, e1)\n",
+    "        \n",
+    "        output = self.dropout(d1)\n",
+    "        return output\n",
+    "    \n",
+    "    \n",
+    "class Critic(nn.Module):\n",
+    "    def __init__(self, in_channels=3):\n",
+    "        super(Critic, self).__init__()\n",
+    "\n",
+    "        def critic_block(in_filters, out_filters, normalization=True):\n",
+    "            layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]\n",
+    "            if normalization:\n",
+    "                layers.append(nn.InstanceNorm2d(out_filters))\n",
+    "            layers.append(nn.LeakyReLU(0.2, inplace=True))\n",
+    "            return layers\n",
+    "\n",
+    "        self.model = nn.Sequential(\n",
+    "            *critic_block(in_channels, 64, normalization=False),\n",
+    "            *critic_block(64, 128),\n",
+    "            *critic_block(128, 256),\n",
+    "            *critic_block(256, 512),\n",
+    "            nn.AdaptiveAvgPool2d(1),\n",
+    "            nn.Flatten(),\n",
+    "            nn.Linear(512, 1)\n",
+    "        )\n",
+    "\n",
+    "    def forward(self, ab, l):\n",
+    "        img_input = torch.cat((ab, l), 1)\n",
+    "        output = self.model(img_input)\n",
+    "        return output\n",
+    "    \n",
+    "    "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from torchsummary import summary\n",
+    "import torch\n",
+    "import os\n",
+    "\n",
+    "class ModelBuilding:\n",
+    "    def __init__(self, config: ModelBuildingConfig):\n",
+    "        self.config = config\n",
+    "        self.root_dir = self.config.root_dir\n",
+    "        self.create_root_dir()\n",
+    "        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "\n",
+    "    def create_root_dir(self):\n",
+    "        os.makedirs(self.root_dir, exist_ok=True)\n",
+    "        print(f\"Created directory: {self.root_dir}\")\n",
+    "\n",
+    "    def get_generator(self):\n",
+    "        return Generator(\n",
+    "            input_channels=self.config.INPUT_CHANNELS,  # corrected argument name\n",
+    "            output_channels=self.config.OUTPUT_CHANNELS,  # corrected argument name\n",
+    "            dropout_rate=self.config.DROPOUT_RATE\n",
+    "        ).to(self.device)\n",
+    "\n",
+    "    def get_critic(self):\n",
+    "        return Critic(in_channels=self.config.IN_CHANNELS).to(self.device)\n",
+    "\n",
+    "    def build(self):\n",
+    "        generator = self.get_generator()\n",
+    "        critic = self.get_critic()\n",
+    "        return generator, critic\n",
+    "\n",
+    "    def save_model(self, model, filename):\n",
+    "        path = self.root_dir / filename\n",
+    "        torch.save(model.state_dict(), path)\n",
+    "        print(f\"Model saved to {path}\")\n",
+    "\n",
+    "    def display_summary(self, model, input_size):\n",
+    "        print(f\"\\nModel Summary:\")\n",
+    "        summary(model, input_size)\n",
+    "\n",
+    "    def build_and_save(self):\n",
+    "        generator, critic = self.build()\n",
+    "\n",
+    "        # Display summaries\n",
+    "        print(\"\\nGenerator Summary:\")\n",
+    "        self.display_summary(generator, (self.config.INPUT_CHANNELS, 224, 224))  # Assuming input size is 224x224\n",
+    "\n",
+    "        print(\"\\nCritic Summary:\")\n",
+    "        self.display_summary(critic, [(2, 224, 224), (1, 224, 224)])  # Critic takes two inputs: ab and l\n",
+    "\n",
+    "        self.save_model(generator, \"generator.pth\")\n",
+    "        self.save_model(critic, \"critic.pth\")\n",
+    "        return generator, critic\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[2024-08-23 00:00:44,340: INFO: common: yaml file: config\\config.yaml loaded successfully]\n",
+      "[2024-08-23 00:00:44,342: INFO: common: yaml file: params.yaml loaded successfully]\n",
+      "[2024-08-23 00:00:44,343: INFO: common: created directory at: artifacts]\n",
+      "Created directory: artifacts\\model\n",
+      "\n",
+      "Generator Summary:\n",
+      "\n",
+      "Model Summary:\n",
+      "----------------------------------------------------------------\n",
+      "        Layer (type)               Output Shape         Param #\n",
+      "================================================================\n",
+      "            Conv2d-1         [-1, 64, 224, 224]             576\n",
+      "       BatchNorm2d-2         [-1, 64, 224, 224]             128\n",
+      "              ReLU-3         [-1, 64, 224, 224]               0\n",
+      "            Conv2d-4         [-1, 64, 224, 224]          36,864\n",
+      "       BatchNorm2d-5         [-1, 64, 224, 224]             128\n",
+      "              ReLU-6         [-1, 64, 224, 224]               0\n",
+      "            Conv2d-7         [-1, 64, 224, 224]             128\n",
+      "              ReLU-8         [-1, 64, 224, 224]               0\n",
+      "          ResBlock-9         [-1, 64, 224, 224]               0\n",
+      "        Dropout2d-10         [-1, 64, 224, 224]               0\n",
+      "        MaxPool2d-11         [-1, 64, 112, 112]               0\n",
+      "           Conv2d-12        [-1, 128, 112, 112]          73,728\n",
+      "      BatchNorm2d-13        [-1, 128, 112, 112]             256\n",
+      "             ReLU-14        [-1, 128, 112, 112]               0\n",
+      "           Conv2d-15        [-1, 128, 112, 112]         147,456\n",
+      "      BatchNorm2d-16        [-1, 128, 112, 112]             256\n",
+      "             ReLU-17        [-1, 128, 112, 112]               0\n",
+      "           Conv2d-18        [-1, 128, 112, 112]           8,320\n",
+      "             ReLU-19        [-1, 128, 112, 112]               0\n",
+      "         ResBlock-20        [-1, 128, 112, 112]               0\n",
+      "   DownsampleConv-21        [-1, 128, 112, 112]               0\n",
+      "        Dropout2d-22        [-1, 128, 112, 112]               0\n",
+      "        MaxPool2d-23          [-1, 128, 56, 56]               0\n",
+      "           Conv2d-24          [-1, 256, 56, 56]         294,912\n",
+      "      BatchNorm2d-25          [-1, 256, 56, 56]             512\n",
+      "             ReLU-26          [-1, 256, 56, 56]               0\n",
+      "           Conv2d-27          [-1, 256, 56, 56]         589,824\n",
+      "      BatchNorm2d-28          [-1, 256, 56, 56]             512\n",
+      "             ReLU-29          [-1, 256, 56, 56]               0\n",
+      "           Conv2d-30          [-1, 256, 56, 56]          33,024\n",
+      "             ReLU-31          [-1, 256, 56, 56]               0\n",
+      "         ResBlock-32          [-1, 256, 56, 56]               0\n",
+      "   DownsampleConv-33          [-1, 256, 56, 56]               0\n",
+      "        Dropout2d-34          [-1, 256, 56, 56]               0\n",
+      "        MaxPool2d-35          [-1, 256, 28, 28]               0\n",
+      "           Conv2d-36          [-1, 512, 28, 28]       1,179,648\n",
+      "      BatchNorm2d-37          [-1, 512, 28, 28]           1,024\n",
+      "             ReLU-38          [-1, 512, 28, 28]               0\n",
+      "           Conv2d-39          [-1, 512, 28, 28]       2,359,296\n",
+      "      BatchNorm2d-40          [-1, 512, 28, 28]           1,024\n",
+      "             ReLU-41          [-1, 512, 28, 28]               0\n",
+      "           Conv2d-42          [-1, 512, 28, 28]         131,584\n",
+      "             ReLU-43          [-1, 512, 28, 28]               0\n",
+      "         ResBlock-44          [-1, 512, 28, 28]               0\n",
+      "   DownsampleConv-45          [-1, 512, 28, 28]               0\n",
+      "        Dropout2d-46          [-1, 512, 28, 28]               0\n",
+      "         Upsample-47          [-1, 512, 56, 56]               0\n",
+      "           Conv2d-48          [-1, 256, 56, 56]       1,769,472\n",
+      "      BatchNorm2d-49          [-1, 256, 56, 56]             512\n",
+      "             ReLU-50          [-1, 256, 56, 56]               0\n",
+      "           Conv2d-51          [-1, 256, 56, 56]         589,824\n",
+      "      BatchNorm2d-52          [-1, 256, 56, 56]             512\n",
+      "             ReLU-53          [-1, 256, 56, 56]               0\n",
+      "           Conv2d-54          [-1, 256, 56, 56]         196,864\n",
+      "             ReLU-55          [-1, 256, 56, 56]               0\n",
+      "         ResBlock-56          [-1, 256, 56, 56]               0\n",
+      "     UpsampleConv-57          [-1, 256, 56, 56]               0\n",
+      "         Upsample-58        [-1, 256, 112, 112]               0\n",
+      "           Conv2d-59        [-1, 128, 112, 112]         442,368\n",
+      "      BatchNorm2d-60        [-1, 128, 112, 112]             256\n",
+      "             ReLU-61        [-1, 128, 112, 112]               0\n",
+      "           Conv2d-62        [-1, 128, 112, 112]         147,456\n",
+      "      BatchNorm2d-63        [-1, 128, 112, 112]             256\n",
+      "             ReLU-64        [-1, 128, 112, 112]               0\n",
+      "           Conv2d-65        [-1, 128, 112, 112]          49,280\n",
+      "             ReLU-66        [-1, 128, 112, 112]               0\n",
+      "         ResBlock-67        [-1, 128, 112, 112]               0\n",
+      "     UpsampleConv-68        [-1, 128, 112, 112]               0\n",
+      "         Upsample-69        [-1, 128, 224, 224]               0\n",
+      "           Conv2d-70         [-1, 64, 224, 224]         110,592\n",
+      "      BatchNorm2d-71         [-1, 64, 224, 224]             128\n",
+      "             ReLU-72         [-1, 64, 224, 224]               0\n",
+      "           Conv2d-73         [-1, 64, 224, 224]          36,864\n",
+      "      BatchNorm2d-74         [-1, 64, 224, 224]             128\n",
+      "             ReLU-75         [-1, 64, 224, 224]               0\n",
+      "           Conv2d-76         [-1, 64, 224, 224]          12,352\n",
+      "             ReLU-77         [-1, 64, 224, 224]               0\n",
+      "         ResBlock-78         [-1, 64, 224, 224]               0\n",
+      "     UpsampleConv-79         [-1, 64, 224, 224]               0\n",
+      "        Dropout2d-80         [-1, 64, 224, 224]               0\n",
+      "================================================================\n",
+      "Total params: 8,216,064\n",
+      "Trainable params: 8,216,064\n",
+      "Non-trainable params: 0\n",
+      "----------------------------------------------------------------\n",
+      "Input size (MB): 0.19\n",
+      "Forward/backward pass size (MB): 1030.53\n",
+      "Params size (MB): 31.34\n",
+      "Estimated Total Size (MB): 1062.06\n",
+      "----------------------------------------------------------------\n",
+      "\n",
+      "Critic Summary:\n",
+      "\n",
+      "Model Summary:\n",
+      "----------------------------------------------------------------\n",
+      "        Layer (type)               Output Shape         Param #\n",
+      "================================================================\n",
+      "            Conv2d-1         [-1, 64, 112, 112]           3,136\n",
+      "         LeakyReLU-2         [-1, 64, 112, 112]               0\n",
+      "            Conv2d-3          [-1, 128, 56, 56]         131,200\n",
+      "    InstanceNorm2d-4          [-1, 128, 56, 56]               0\n",
+      "         LeakyReLU-5          [-1, 128, 56, 56]               0\n",
+      "            Conv2d-6          [-1, 256, 28, 28]         524,544\n",
+      "    InstanceNorm2d-7          [-1, 256, 28, 28]               0\n",
+      "         LeakyReLU-8          [-1, 256, 28, 28]               0\n",
+      "            Conv2d-9          [-1, 512, 14, 14]       2,097,664\n",
+      "   InstanceNorm2d-10          [-1, 512, 14, 14]               0\n",
+      "        LeakyReLU-11          [-1, 512, 14, 14]               0\n",
+      "AdaptiveAvgPool2d-12            [-1, 512, 1, 1]               0\n",
+      "          Flatten-13                  [-1, 512]               0\n",
+      "           Linear-14                    [-1, 1]             513\n",
+      "================================================================\n",
+      "Total params: 2,757,057\n",
+      "Trainable params: 2,757,057\n",
+      "Non-trainable params: 0\n",
+      "----------------------------------------------------------------\n",
+      "Input size (MB): 2824.00\n",
+      "Forward/backward pass size (MB): 28.34\n",
+      "Params size (MB): 10.52\n",
+      "Estimated Total Size (MB): 2862.85\n",
+      "----------------------------------------------------------------\n",
+      "Model saved to artifacts\\model\\generator.pth\n",
+      "Model saved to artifacts\\model\\critic.pth\n"
+     ]
+    }
+   ],
+   "source": [
+    "try:\n",
+    "    config_manager = ConfigurationManager()\n",
+    "    model_config = config_manager.get_model_building_config()\n",
+    "\n",
+    "    model_building = ModelBuilding(config=model_config)\n",
+    "    generator, critic = model_building.build_and_save()\n",
+    "except Exception as e:\n",
+    "    raise e"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

src/imagecolorization/config/configuration.py

@@ -1,7 +1,8 @@
 from src.imagecolorization.constants import *
 from src.imagecolorization.utils.common import read_yaml, create_directories
 from src.imagecolorization.entity.config_entity import (DataIngestionConfig,
-                                                        DataTransformationConfig)
+                                                        DataTransformationConfig,
+                                                        ModelBuildingConfig)
 class ConfigurationManager:
     def __init__(
         self,
@@ -44,6 +45,26 @@ class ConfigurationManager:
         )
         
         return data_transformation_cofig
+    
+    def get_model_building_config(self) -> ModelBuildingConfig:
+        config = self.config.model_building
+        params = self.params
+
+        model_building_config = ModelBuildingConfig(
+            root_dir=Path(config.root_dir),
+            KERNEL_SIZE_RES=params.KERNEL_SIZE_RES,
+            PADDING=params.PADDING,
+            STRIDE=params.STRIDE,
+            BIAS=params.BIAS,
+            SCALE_FACTOR=params.SCALE_FACTOR,
+            DIM=params.DIM,
+            DROPOUT_RATE=params.DROPOUT_RATE,
+            KERNEL_SIZE_GENERATOR=params.KERNEL_SIZE_GENERATOR,
+            INPUT_CHANNELS=params.INPUT_CHANNELS,
+            OUTPUT_CHANNELS=params.OUTPUT_CHANNELS,
+            IN_CHANNELS=params.IN_CHANNELS
+        )
+        return model_building_config
         
     
     

src/imagecolorization/conponents/model_building.py

@@ -0,0 +1,168 @@
+import torch 
+import torch.nn as nn
+from pathlib import Path
+from torchsummary import summary
+import os
+from src.imagecolorization.config.configuration import ModelBuildingConfig
+
+class ResBlock(nn.Module):
+    def __init__(self, in_channles, out_channels, stride = 1, kerenl_size = 3, padding = 1, bias = False):
+        super().__init__()
+        self.layer = nn.Sequential(
+            nn.Conv2d(in_channles, out_channels, kernel_size=kerenl_size, padding=padding, stride=stride, bias = bias),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=kerenl_size, padding=padding, stride = 1, bias = bias),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+        
+        self.identity_map = nn.Conv2d(in_channles, out_channels,kernel_size=1, stride=stride)
+        self.relu = nn.ReLU(inplace= True)
+        
+    def forward(self, inputs):
+        x = inputs.clone().detach()
+        out = self.layer(x)
+        residual = self.identity_map(inputs)
+        skip  = out + residual
+        return self.relu(skip)
+    
+    
+class DownsampleConv(nn.Module):
+    def __init__(self, in_channels, out_channels, stride = 1):
+        super().__init__()
+        self.layer = nn.Sequential(
+            nn.MaxPool2d(2),
+            ResBlock(in_channels, out_channels)
+        )
+        
+    def forward(self, inputs):
+        return self.layer(inputs)
+    
+    
+    
+class UpsampleConv(nn.Module):
+    def __init__(self, in_channels, out_channels, scale_factor=2):
+        super().__init__()
+        self.upsample = nn.Upsample(scale_factor=scale_factor,mode = 'bilinear', align_corners=True)
+        self.res_block = ResBlock(in_channels + out_channels, out_channels)
+
+    def forward(self, inputs, skip):
+        x = self.upsample(inputs)
+        x = torch.cat([x, skip], dim = 1)
+        x = self.res_block(x)
+        return x
+    
+class Generator(nn.Module):
+    def __init__(self, input_channels, output_channels, dropout_rate = 0.2):
+        super().__init__()
+        self.encoding_layer1_= ResBlock(input_channels, 64)
+        self.encoding_layer2_ = DownsampleConv(64, 128)
+        self.encoding_layer3_ = DownsampleConv(128, 256)
+        self.bridge = DownsampleConv(256, 512)
+        self.decoding_layer3 = UpsampleConv(512, 256)
+        self.decoding_layer2 = UpsampleConv(256, 128)
+        self.decoding_layer1 = UpsampleConv(128 , 64)
+        self.output = nn.Conv2d(64, output_channels, kernel_size = 1)
+        self.dropout = nn.Dropout2d(dropout_rate)
+        
+    def forward(self, inputs):
+        e1 = self.encoding_layer1_(inputs)
+        e1 = self.dropout(e1)
+        e2 = self.encoding_layer2_(e1)
+        e2 = self.dropout(e2)
+        e3 = self.encoding_layer3_(e2)
+        e3 = self.dropout(e3)
+        
+        bridge = self.bridge(e3)
+        bridge = self.dropout(bridge)
+        
+        d3 = self.decoding_layer3(bridge, e3)
+        d2 =self.decoding_layer2(d3, e2)
+        d1 = self.decoding_layer1(d2, e1)
+        
+        output = self.dropout(d1)
+        return output
+    
+    
+class Critic(nn.Module):
+    def __init__(self, in_channels=3):
+        super(Critic, self).__init__()
+
+        def critic_block(in_filters, out_filters, normalization=True):
+            layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]
+            if normalization:
+                layers.append(nn.InstanceNorm2d(out_filters))
+            layers.append(nn.LeakyReLU(0.2, inplace=True))
+            return layers
+
+        self.model = nn.Sequential(
+            *critic_block(in_channels, 64, normalization=False),
+            *critic_block(64, 128),
+            *critic_block(128, 256),
+            *critic_block(256, 512),
+            nn.AdaptiveAvgPool2d(1),
+            nn.Flatten(),
+            nn.Linear(512, 1)
+        )
+
+    def forward(self, ab, l):
+        img_input = torch.cat((ab, l), 1)
+        output = self.model(img_input)
+        return output
+    
+    
+
+
+class ModelBuilding:
+    def __init__(self, config: ModelBuildingConfig):
+        self.config = config
+        self.root_dir = self.config.root_dir
+        self.create_root_dir()
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    def create_root_dir(self):
+        os.makedirs(self.root_dir, exist_ok=True)
+        print(f"Created directory: {self.root_dir}")
+
+    def get_generator(self):
+        return Generator(
+            input_channels=self.config.INPUT_CHANNELS,  # corrected argument name
+            output_channels=self.config.OUTPUT_CHANNELS,  # corrected argument name
+            dropout_rate=self.config.DROPOUT_RATE
+        ).to(self.device)
+
+    def get_critic(self):
+        return Critic(in_channels=self.config.IN_CHANNELS).to(self.device)
+
+    def build(self):
+        generator = self.get_generator()
+        critic = self.get_critic()
+        return generator, critic
+
+    def save_model(self, model, filename):
+        path = self.root_dir / filename
+        torch.save(model.state_dict(), path)
+        print(f"Model saved to {path}")
+
+    def display_summary(self, model, input_size):
+        print(f"\nModel Summary:")
+        summary(model, input_size)
+
+    def build_and_save(self):
+        generator, critic = self.build()
+
+        # Display summaries
+        print("\nGenerator Summary:")
+        self.display_summary(generator, (self.config.INPUT_CHANNELS, 224, 224))  # Assuming input size is 224x224
+
+        print("\nCritic Summary:")
+        self.display_summary(critic, [(2, 224, 224), (1, 224, 224)])  # Critic takes two inputs: ab and l
+
+        self.save_model(generator, "generator.pth")
+        self.save_model(critic, "critic.pth")
+        return generator, critic
+
+
+    
+    

src/imagecolorization/entity/config_entity.py

@@ -16,4 +16,20 @@ class DataTransformationConfig:
     data_path_grey : Path
     BATCH_SIZE : int
     IMAGE_SIZE : list
-    DATA_RANGE: int
+    DATA_RANGE: int
+    
+    
+@dataclass(frozen=True)
+class ModelBuildingConfig:
+    root_dir: Path
+    KERNEL_SIZE_RES: int
+    PADDING: int
+    STRIDE: int
+    BIAS: bool
+    SCALE_FACTOR: int
+    DIM: int
+    DROPOUT_RATE: float
+    KERNEL_SIZE_GENERATOR: int
+    INPUT_CHANNELS: int
+    OUTPUT_CHANNELS: int
+    IN_CHANNELS: int

src/imagecolorization/pipeline/stage_03_model_building.py

@@ -0,0 +1,13 @@
+from src.imagecolorization.conponents.model_building import ModelBuilding
+from src.imagecolorization.config.configuration import ConfigurationManager
+
+class ModelBuildingPipeline:
+    def __init__(slef):
+        pass
+    
+    def main(self):
+        config_manager = ConfigurationManager()
+        model_config = config_manager.get_model_building_config()
+
+        model_building = ModelBuilding(config=model_config)
+        generator, critic = model_building.build_and_save()