commiting backend files

.gitignore

@@ -1,207 +1,105 @@
-# Byte-compiled / optimized / DLL files
+# =============================================================================
+# Python Virtual Environments
+# Never commit the virtual environment folder.
+# =============================================================================
+/venv/
+.venv/
+env/
+.env
+ENV/
+env.bak/
+venv.bak/
+
+
+# =============================================================================
+# DVC and MLflow Artifacts
+# DVC tracks data, so we don't need Git to. We only commit the .dvc files.
+# The `artifacts` and `logs` directories will be generated by the pipeline.
+# =============================================================================
+/artifacts/
+/logs/
+/mlruns/
+# DVC's internal cache should NEVER be committed.
+/.dvc/cache/
+# DVC's temporary directories
+/.dvc/tmp/
+
+
+# =============================================================================
+# Python Cache and Compiled Files
+# These are generated automatically by Python and don't need to be versioned.
+# =============================================================================
 __pycache__/
-*.py[codz]
+*.py[cod]
 *$py.class
 
-# C extensions
-*.so
-
-# Distribution / packaging
-.Python
-build/
-develop-eggs/
-dist/
-downloads/
-eggs/
-.eggs/
-lib/
-lib64/
-parts/
-sdist/
-var/
-wheels/
-share/python-wheels/
+
+# =============================================================================
+# Build and Distribution Artifacts
+# Generated when building a Python package.
+# =============================================================================
+/build/
+/develop-eggs/
+/dist/
+/downloads/
+/eggs/
+/.eggs/
+/lib/
+/lib64/
+/parts/
+/sdist/
+/var/
+/wheels/
 *.egg-info/
 .installed.cfg
 *.egg
-MANIFEST
-
-# PyInstaller
-#  Usually these files are written by a python script from a template
-#  before PyInstaller builds the exe, so as to inject date/other infos into it.
 *.manifest
 *.spec
 
-# Installer logs
-pip-log.txt
-pip-delete-this-directory.txt
-
-# Unit test / coverage reports
-htmlcov/
-.tox/
-.nox/
-.coverage
-.coverage.*
-.cache
-nosetests.xml
-coverage.xml
-*.cover
-*.py.cover
-.hypothesis/
-.pytest_cache/
-cover/
-
-# Translations
-*.mo
-*.pot
-
-# Django stuff:
-*.log
-local_settings.py
-db.sqlite3
-db.sqlite3-journal
-
-# Flask stuff:
-instance/
-.webassets-cache
-
-# Scrapy stuff:
-.scrapy
-
-# Sphinx documentation
-docs/_build/
-
-# PyBuilder
-.pybuilder/
-target/
-
-# Jupyter Notebook
-.ipynb_checkpoints
-
-# IPython
-profile_default/
-ipython_config.py
-
-# pyenv
-#   For a library or package, you might want to ignore these files since the code is
-#   intended to run in multiple environments; otherwise, check them in:
-# .python-version
-
-# pipenv
-#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
-#   However, in case of collaboration, if having platform-specific dependencies or dependencies
-#   having no cross-platform support, pipenv may install dependencies that don't work, or not
-#   install all needed dependencies.
-#Pipfile.lock
-
-# UV
-#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
-#   This is especially recommended for binary packages to ensure reproducibility, and is more
-#   commonly ignored for libraries.
-#uv.lock
-
-# poetry
-#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
-#   This is especially recommended for binary packages to ensure reproducibility, and is more
-#   commonly ignored for libraries.
-#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
-#poetry.lock
-#poetry.toml
-
-# pdm
-#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
-#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
-#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
-#pdm.lock
-#pdm.toml
-.pdm-python
-.pdm-build/
-
-# pixi
-#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
-#pixi.lock
-#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
-#   in the .venv directory. It is recommended not to include this directory in version control.
-.pixi
-
-# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
-__pypackages__/
-
-# Celery stuff
-celerybeat-schedule
-celerybeat.pid
-
-# SageMath parsed files
-*.sage.py
-
-# Environments
-.env
-.envrc
-.venv
-env/
-venv/
-ENV/
-env.bak/
-venv.bak/
 
-# Spyder project settings
-.spyderproject
-.spyproject
-
-# Rope project settings
-.ropeproject
-
-# mkdocs documentation
-/site
-
-# mypy
-.mypy_cache/
-.dmypy.json
-dmypy.json
-
-# Pyre type checker
-.pyre/
-
-# pytype static type analyzer
-.pytype/
-
-# Cython debug symbols
-cython_debug/
-
-# PyCharm
-#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
-#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
-#  and can be added to the global gitignore or merged into this file.  For a more nuclear
-#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
-
-# Abstra
-# Abstra is an AI-powered process automation framework.
-# Ignore directories containing user credentials, local state, and settings.
-# Learn more at https://abstra.io/docs
-.abstra/
-
-# Visual Studio Code
-#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
-#  that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
-#  and can be added to the global gitignore or merged into this file. However, if you prefer, 
-#  you could uncomment the following to ignore the entire vscode folder
-# .vscode/
-
-# Ruff stuff:
-.ruff_cache/
-
-# PyPI configuration file
-.pypirc
-
-# Cursor
-#  Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
-#  exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
-#  refer to https://docs.cursor.com/context/ignore-files
-.cursorignore
-.cursorindexingignore
-
-# Marimo
-marimo/_static/
-marimo/_lsp/
-__marimo__/
+# =============================================================================
+# IDE and Editor-Specific Files
+# Ignore configuration files from common editors like VSCode, PyCharm, etc.
+# =============================================================================
+.vscode/
+.idea/
+.project
+.pydevproject
+.classpath
+*.swp
+*.swo
+
+
+# =============================================================================
+# OS-specific Files
+# Ignore files generated by macOS, Windows, and Linux.
+# =============================================================================
+.DS_Store
+Thumbs.db
+Desktop.ini
+
+
+# =============================================================================
+# Jupyter Notebook Checkpoints
+# Ignore the checkpoints directory created by Jupyter.
+# =============================================================================
+.ipynb_checkpoints/
+
+
+# =============================================================================
+# Kaggle API Credentials
+# IMPORTANT: Never commit your API keys or secrets.
+# =============================================================================
+kaggle.json
+
+
+# =============================================================================
+# Other
+# Any other miscellaneous files that shouldn't be in the repo.
+# =============================================================================
+*.log
+*.tmp
+*.bak
+*.local
+eval_output/ # Temporary directory created by the evaluation component
+.env

config/config.yaml

@@ -0,0 +1,39 @@
+# config/config.yaml
+
+artifacts_root: artifacts
+
+data_ingestion:
+  root_dir: artifacts/data_ingestion
+  source_kaggle_dataset_id: "paultimothymooney/chest-xray-pneumonia"
+  unzip_dir: artifacts/data_ingestion/
+  # We will create these three files now
+  train_df_path: artifacts/data_ingestion/train_df.csv
+  test_df_path: artifacts/data_ingestion/test_df.csv
+  val_df_path: artifacts/data_ingestion/val_df.csv
+
+data_transformation:
+  root_dir: artifacts/data_transformation
+  # We now have three sources
+  train_data_path: artifacts/data_ingestion/train_df.csv
+  test_data_path: artifacts/data_ingestion/test_df.csv
+  val_data_path: artifacts/data_ingestion/val_df.csv
+  # And will create three outputs
+  train_dataset_path: artifacts/data_transformation/train_dataset
+  test_dataset_path: artifacts/data_transformation/test_dataset
+  val_dataset_path: artifacts/data_transformation/val_dataset
+
+model_training:
+  root_dir: artifacts/model_training
+  trained_model_path: artifacts/model_training/model
+  model_name: "google/vit-base-patch16-224-in21k"
+  # We'll use the validation set for evaluation during training
+  train_dataset_path: artifacts/data_transformation/train_dataset
+  val_dataset_path: artifacts/data_transformation/val_dataset
+
+model_evaluation:
+  root_dir: artifacts/model_evaluation
+  model_path: artifacts/model_training/model
+  # Final evaluation is done on the unseen test set
+  test_dataset_path: artifacts/data_transformation/test_dataset
+  metrics_file_name: artifacts/model_evaluation/metrics.json
+  mlflow_uri: "https://dagshub.com/AlyyanAhmed21/Chest-X-ray-Pneumonia-Detection-with-ViT.mlflow"

dvc.lock

@@ -0,0 +1,109 @@
+schema: '2.0'
+stages:
+  data_ingestion:
+    cmd: python src/vitClassifier/pipeline/stage_01_data_ingestion.py
+    deps:
+    - path: config/config.yaml
+      hash: md5
+      md5: 9a45c00c11b9a8a0b4e396952a6b19a8
+      size: 1591
+    - path: src/vitClassifier/pipeline/stage_01_data_ingestion.py
+      hash: md5
+      md5: 2486829b866dffe25a752248afa95f4f
+      size: 968
+    outs:
+    - path: artifacts/data_ingestion
+      hash: md5
+      md5: 9921508de1b9f2e8a5a4150d063e178d.dir
+      size: 2484429974
+      nfiles: 17594
+  data_transformation:
+    cmd: python src/vitClassifier/pipeline/stage_02_data_transformation.py
+    deps:
+    - path: artifacts/data_ingestion/test_df.csv
+      hash: md5
+      md5: 95cbf91a4d0719e528c74879d6da0e34
+      size: 53272
+    - path: artifacts/data_ingestion/train_df.csv
+      hash: md5
+      md5: af0d24afd4d9092b64bb1db986d38f76
+      size: 460017
+    - path: artifacts/data_ingestion/val_df.csv
+      hash: md5
+      md5: 575134cde7f8113c2a51dd4fac3e4c5e
+      size: 1389
+    - path: config/config.yaml
+      hash: md5
+      md5: 9a45c00c11b9a8a0b4e396952a6b19a8
+      size: 1591
+    - path: params.yaml
+      hash: md5
+      md5: cc525f2481819601bb93ec5d7f008dda
+      size: 127
+    - path: src/vitClassifier/pipeline/stage_02_data_transformation.py
+      hash: md5
+      md5: 095fcfa8843d6b94a05d9f2172522b32
+      size: 1237
+    outs:
+    - path: artifacts/data_transformation
+      hash: md5
+      md5: b91edf7d0a4b4b0022f2d33d3f2176fa.dir
+      size: 5074488112
+      nfiles: 18
+  model_training:
+    cmd: python src/vitClassifier/pipeline/stage_03_model_training.py
+    deps:
+    - path: artifacts/data_transformation/train_dataset
+      hash: md5
+      md5: 64425f8e57c16ac250c2ea73b78b7aa2.dir
+      size: 4687397188
+      nfiles: 12
+    - path: artifacts/data_transformation/val_dataset
+      hash: md5
+      md5: e19e673f104f2efec21df351b3d4869c.dir
+      size: 9678966
+      nfiles: 3
+    - path: config/config.yaml
+      hash: md5
+      md5: 9a45c00c11b9a8a0b4e396952a6b19a8
+      size: 1591
+    - path: params.yaml
+      hash: md5
+      md5: cc525f2481819601bb93ec5d7f008dda
+      size: 127
+    - path: src/vitClassifier/pipeline/stage_03_model_training.py
+      hash: md5
+      md5: 5e9cde1828fc4b2608e9f4e92b134a07
+      size: 815
+    outs:
+    - path: artifacts/model_training/model
+      hash: md5
+      md5: 9f0765ff59616eddac47fcaf7a5e7387.dir
+      size: 343230531
+      nfiles: 4
+  model_evaluation:
+    cmd: python src/vitClassifier/pipeline/stage_04_model_evaluation.py
+    deps:
+    - path: artifacts/data_transformation/test_dataset
+      hash: md5
+      md5: 41a8f95d5075f06bef31fbf55d838cca.dir
+      size: 377411958
+      nfiles: 3
+    - path: artifacts/model_training/model
+      hash: md5
+      md5: 9f0765ff59616eddac47fcaf7a5e7387.dir
+      size: 343230531
+      nfiles: 4
+    - path: config/config.yaml
+      hash: md5
+      md5: 9224d2383ec670f1738b47139f250ad4
+      size: 1659
+    - path: src/vitClassifier/pipeline/stage_04_model_evaluation.py
+      hash: md5
+      md5: e31c602e23dbfa62f6453ca44b621d0a
+      size: 863
+    outs:
+    - path: artifacts/model_evaluation/metrics.json
+      hash: md5
+      md5: 26b4e3326f589929e4a6e34833cc187f
+      size: 150

dvc.yaml

@@ -0,0 +1,44 @@
+stages:
+  data_ingestion:
+    cmd: python src/vitClassifier/pipeline/stage_01_data_ingestion.py
+    deps:
+      - src/vitClassifier/pipeline/stage_01_data_ingestion.py
+      - config/config.yaml
+    outs:
+      - artifacts/data_ingestion
+
+  data_transformation:
+    cmd: python src/vitClassifier/pipeline/stage_02_data_transformation.py
+    deps:
+      - src/vitClassifier/pipeline/stage_02_data_transformation.py
+      # --- THIS IS THE FIX ---
+      # Remove the old dependency and add the three new ones.
+      - artifacts/data_ingestion/train_df.csv
+      - artifacts/data_ingestion/test_df.csv
+      - artifacts/data_ingestion/val_df.csv
+      - config/config.yaml
+      - params.yaml
+    outs:
+      - artifacts/data_transformation
+
+  model_training:
+    cmd: python src/vitClassifier/pipeline/stage_03_model_training.py
+    deps:
+      - src/vitClassifier/pipeline/stage_03_model_training.py
+      - artifacts/data_transformation/train_dataset
+      - artifacts/data_transformation/val_dataset # Added dependency on val dataset
+      - config/config.yaml
+      - params.yaml
+    outs:
+      - artifacts/model_training/model
+
+  model_evaluation:
+    cmd: python src/vitClassifier/pipeline/stage_04_model_evaluation.py
+    deps:
+      - src/vitClassifier/pipeline/stage_04_model_evaluation.py
+      - artifacts/data_transformation/test_dataset
+      - artifacts/model_training/model
+      - config/config.yaml
+    metrics:
+    - artifacts/model_evaluation/metrics.json:
+        cache: false

gpu

@@ -0,0 +1,57 @@
+# check_gpu.py
+
+import sys
+import torch
+
+def check_gpu_environment():
+    """
+    This script checks the system's Python and PyTorch GPU environment.
+    It prints detailed information about the setup.
+    """
+    print("--- System and Python Information ---")
+    print(f"Python Version: {sys.version}")
+    print("\n--- PyTorch and CUDA Information ---")
+    
+    try:
+        print(f"PyTorch Version: {torch.__version__}")
+        
+        # Check if CUDA (GPU support) is available
+        cuda_available = torch.cuda.is_available()
+        print(f"CUDA Available: {cuda_available}")
+        
+        if not cuda_available:
+            print("\nWARNING: PyTorch was not built with CUDA support. GPU will not be used.")
+            return
+
+        # Get the number of available GPUs
+        gpu_count = torch.cuda.device_count()
+        print(f"Number of GPUs Available: {gpu_count}")
+
+        # Get details for each GPU
+        for i in range(gpu_count):
+            print(f"\n--- GPU Details (Device {i}) ---")
+            gpu_name = torch.cuda.get_device_name(i)
+            print(f"  GPU Name: {gpu_name}")
+            
+            cuda_capability = torch.cuda.get_device_capability(i)
+            print(f"  Compute Capability: {cuda_capability[0]}.{cuda_capability[1]}")
+            
+            total_mem = torch.cuda.get_device_properties(i).total_memory / (1024**3) # Convert bytes to GB
+            print(f"  Total Memory: {total_mem:.2f} GB")
+
+        # Check for cuDNN
+        cudnn_available = torch.backends.cudnn.is_available()
+        print("\n--- cuDNN Information ---")
+        print(f"cuDNN Available: {cudnn_available}")
+        if cudnn_available:
+            cudnn_version = torch.backends.cudnn.version()
+            print(f"cuDNN Version: {cudnn_version}")
+        else:
+            print("\nWARNING: cuDNN is not available. Training will be significantly slower.")
+
+    except Exception as e:
+        print(f"\nAn error occurred: {e}")
+        print("Please ensure PyTorch is installed correctly.")
+
+if __name__ == "__main__":
+    check_gpu_environment()

main.py

@@ -0,0 +1,23 @@
+from vitClassifier import logger
+from vitClassifier.pipeline.stage_01_data_ingestion import DataIngestionTrainingPipeline
+from vitClassifier.pipeline.stage_02_data_transformation import DataTransformationTrainingPipeline
+from vitClassifier.pipeline.stage_03_model_training import ModelTrainingPipeline
+from vitClassifier.pipeline.stage_04_model_evaluation import ModelEvaluationPipeline
+from dotenv import load_dotenv
+load_dotenv()
+
+def run_pipeline(stage_name, pipeline_class):
+    try:
+        logger.info(f">>>>>> stage {stage_name} started <<<<<<")
+        pipeline = pipeline_class()
+        pipeline.main()
+        logger.info(f">>>>>> stage {stage_name} completed <<<<<<\n\nx==========x")
+    except Exception as e:
+        logger.exception(e)
+        raise e
+
+if __name__ == '__main__':
+    run_pipeline("Data Ingestion stage", DataIngestionTrainingPipeline)
+    run_pipeline("Data Transformation stage", DataTransformationTrainingPipeline)
+    run_pipeline("Model Training stage", ModelTrainingPipeline)
+    run_pipeline("Model Evaluation stage", ModelEvaluationPipeline)

params.yaml

@@ -0,0 +1,7 @@
+LEARNING_RATE: 2.0e-5
+BATCH_SIZE: 32
+EPOCHS: 3
+WEIGHT_DECAY: 0.01
+WARMUP_STEPS: 100
+RANDOM_STATE: 42
+TEST_SPLIT_SIZE: 0.2

requirements.txt

@@ -0,0 +1,21 @@
+pandas
+numpy
+torch
+torchvision
+transformers
+datasets>=2.14.5
+evaluate
+accelerate>=0.27
+mlflow
+scikit-learn
+imblearn
+python-box
+PyYAML
+ensure
+tqdm
+pathlib
+dvc
+matplotlib
+Pillow
+kaggle
+python-dotenv

setup.py

@@ -0,0 +1,27 @@
+import setuptools
+
+with open("README.md", "r", encoding="utf-8") as f:
+    long_description = f.read()
+
+__version__ = "0.0.1"
+
+REPO_NAME = "Chest-Xray-Pneumonia-ViT-MLflow-DVC"
+AUTHOR_USER_NAME = "your-github-username" # CHANGE THIS
+SRC_REPO = "vitClassifier"
+AUTHOR_EMAIL = "your-email@example.com" # CHANGE THIS
+
+setuptools.setup(
+    name=SRC_REPO,
+    version=__version__,
+    author=AUTHOR_USER_NAME,
+    author_email=AUTHOR_EMAIL,
+    description="An end-to-end ML project for Chest X-ray Pneumonia classification using ViT.",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url=f"https://github.com/{AUTHOR_USER_NAME}/{REPO_NAME}",
+    project_urls={
+        "Bug Tracker": f"https://github.com/{AUTHOR_USER_NAME}/{REPO_NAME}/issues",
+    },
+    package_dir={"": "src"},
+    packages=setuptools.find_packages(where="src")
+)

src/vitClassifier/__init__.py

@@ -0,0 +1,26 @@
+# src/vitClassifier/__init__.py
+
+import os
+import sys
+import logging
+
+# Define the logging format
+logging_str = "[%(asctime)s: %(levelname)s: %(module)s: %(message)s]"
+
+# Define the directory for log files
+log_dir = "logs"
+log_filepath = os.path.join(log_dir, "running_logs.log")
+os.makedirs(log_dir, exist_ok=True)
+
+# Configure the logging
+logging.basicConfig(
+    level=logging.INFO,
+    format=logging_str,
+    handlers=[
+        logging.FileHandler(log_filepath),  # Log to a file
+        logging.StreamHandler(sys.stdout)   # Also log to the console
+    ]
+)
+
+# Create a logger object that can be imported by other modules
+logger = logging.getLogger("vitClassifierLogger")

src/vitClassifier/components/data_ingestion.py

@@ -0,0 +1,67 @@
+import os
+import pandas as pd
+from pathlib import Path
+from vitClassifier import logger
+from vitClassifier.entity.config_entity import DataIngestionConfig
+import kaggle
+
+class DataIngestion:
+    def __init__(self, config: DataIngestionConfig):
+        self.config = config
+
+    def download_dataset(self):
+        try:
+            # ... (download logic remains exactly the same)
+            logger.info("Authenticating with Kaggle API...")
+            kaggle.api.authenticate()
+            logger.info("Authentication successful.")
+
+            dataset_id = self.config.source_kaggle_dataset_id
+            download_path = self.config.unzip_dir
+            
+            expected_data_folder = download_path / "chest_xray"
+            if expected_data_folder.exists():
+                logger.info(f"Dataset already exists at {expected_data_folder}. Skipping download.")
+                return
+
+            logger.info(f"Downloading dataset '{dataset_id}' to '{download_path}'...")
+            kaggle.api.dataset_download_files(
+                dataset=dataset_id, path=download_path, unzip=True, quiet=False
+            )
+            logger.info("Dataset downloaded and unzipped successfully.")
+
+        except Exception as e:
+            logger.error(f"Failed to download dataset from Kaggle. Error: {e}")
+            raise e
+
+    def create_dataframes(self):
+        """
+        Scans train, test, and val directories and creates separate DataFrames.
+        """
+        source_root = self.config.unzip_dir / "chest_xray"
+        
+        # Helper function to create a dataframe for a given split (train/test/val)
+        def _create_df_for_split(split_name: str, save_path: Path):
+            split_path = source_root / split_name
+            file_names, labels = [], []
+            
+            # Using .glob to find all .jpeg files in NORMAL and PNEUMONIA subfolders
+            for file in sorted(split_path.glob('*/*.jpeg')):
+                label = file.parent.name # NORMAL or PNEUMONIA
+                labels.append(label)
+                file_names.append(str(file))
+            
+            df = pd.DataFrame({"image": file_names, "label": labels})
+            df.to_csv(save_path, index=False)
+            logger.info(f"Created and saved {split_name} DataFrame to {save_path}")
+
+        # Create DataFrames for each split
+        _create_df_for_split("train", self.config.train_df_path)
+        _create_df_for_split("test", self.config.test_df_path)
+        _create_df_for_split("val", self.config.val_df_path)
+
+    def ingest_data(self):
+        logger.info("Starting data ingestion process.")
+        self.download_dataset()
+        self.create_dataframes()
+        logger.info("Data ingestion process completed.")

src/vitClassifier/components/data_transformation.py

@@ -0,0 +1,80 @@
+# src/vitClassifier/components/data_transformation.py
+
+import pandas as pd
+from datasets import Dataset, Image, ClassLabel
+from imblearn.over_sampling import RandomOverSampler
+from vitClassifier.entity.config_entity import DataTransformationConfig
+from vitClassifier import logger
+# --- NEW IMPORTS ---
+from transformers import ViTImageProcessor
+from torchvision.transforms import (Compose, Resize, ToTensor, Normalize, RandomRotation, RandomHorizontalFlip)
+
+class DataTransformation:
+    def __init__(self, config: DataTransformationConfig, random_state: int, model_name: str):
+        self.config = config
+        self.random_state = random_state
+        self.model_name = model_name # <-- Need model_name to load the correct processor
+
+    def transform_data(self):
+        # --- 1. Load DataFrames and apply Oversampling (same as before) ---
+        train_df = pd.read_csv(self.config.train_data_path)
+        test_df = pd.read_csv(self.config.test_data_path)
+        val_df = pd.read_csv(self.config.val_data_path)
+        
+        y = train_df[['label']]
+        X = train_df.drop(['label'], axis=1)
+        ros = RandomOverSampler(random_state=self.random_state)
+        X_resampled, y_resampled = ros.fit_resample(X, y)
+        train_df_balanced = pd.concat([X_resampled, y_resampled], axis=1)
+        
+        train_dataset = Dataset.from_pandas(train_df_balanced).cast_column("image", Image())
+        test_dataset = Dataset.from_pandas(test_df).cast_column("image", Image())
+        val_dataset = Dataset.from_pandas(val_df).cast_column("image", Image())
+        
+        # --- 2. Label Encoding (same as before) ---
+        labels_list = train_df_balanced['label'].unique().tolist()
+        class_labels = ClassLabel(num_classes=len(labels_list), names=labels_list)
+
+        def map_label2id(example):
+            example['label'] = class_labels.str2int(example['label'])
+            return example
+
+        train_dataset = train_dataset.map(map_label2id, batched=True).cast_column('label', class_labels)
+        test_dataset = test_dataset.map(map_label2id, batched=True).cast_column('label', class_labels)
+        val_dataset = val_dataset.map(map_label2id, batched=True).cast_column('label', class_labels)
+
+        # --- 3. THE NEW LOGIC: Preprocess images with .map() ---
+        logger.info("Starting image preprocessing with .map(). This may take a few minutes...")
+        processor = ViTImageProcessor.from_pretrained(self.model_name)
+        image_mean, image_std = processor.image_mean, processor.image_std
+        size = processor.size["height"]
+        normalize = Normalize(mean=image_mean, std=image_std)
+
+        # Define transforms
+        _train_transforms = Compose([Resize((size, size)), RandomRotation(15), RandomHorizontalFlip(), ToTensor(), normalize])
+        _val_test_transforms = Compose([Resize((size, size)), ToTensor(), normalize])
+        
+        def apply_train_transforms(examples):
+            examples['pixel_values'] = [_train_transforms(image.convert("RGB")) for image in examples['image']]
+            return examples
+
+        def apply_val_test_transforms(examples):
+            examples['pixel_values'] = [_val_test_transforms(image.convert("RGB")) for image in examples['image']]
+            return examples
+            
+        # Use .map() to apply transforms and create 'pixel_values' column
+        train_dataset = train_dataset.map(apply_train_transforms, batched=True)
+        test_dataset = test_dataset.map(apply_val_test_transforms, batched=True)
+        val_dataset = val_dataset.map(apply_val_test_transforms, batched=True)
+
+        # Remove the original 'image' column to save space
+        train_dataset = train_dataset.remove_columns(['image'])
+        test_dataset = test_dataset.remove_columns(['image'])
+        val_dataset = val_dataset.remove_columns(['image'])
+
+        # --- 4. Save the fully processed datasets ---
+        train_dataset.save_to_disk(str(self.config.train_dataset_path))
+        test_dataset.save_to_disk(str(self.config.test_dataset_path))
+        val_dataset.save_to_disk(str(self.config.val_dataset_path))
+        
+        logger.info("Data Transformation complete. Fully preprocessed datasets saved.")

src/vitClassifier/components/model_evaluation.py

@@ -0,0 +1,89 @@
+# src/vitClassifier/components/model_evaluation.py
+
+import mlflow
+import mlflow.pytorch
+import torch
+import json
+from pathlib import Path
+from datasets import load_from_disk
+from transformers import (ViTForImageClassification, ViTImageProcessor, Trainer, TrainingArguments, DefaultDataCollator)
+from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
+from vitClassifier.entity.config_entity import EvaluationConfig
+from vitClassifier.utils.common import read_yaml # Keep this if you need it, but it's not used here
+from vitClassifier import logger
+
+class ModelEvaluation:
+    def __init__(self, config: EvaluationConfig):
+        self.config = config
+
+    def evaluate(self):
+        # Determine device
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        # Load the best model from the training stage and move it to the correct device
+        model_path = str(self.config.path_of_model)
+        model = ViTForImageClassification.from_pretrained(model_path).to(device)
+        
+        # Load the pre-processed test dataset
+        test_data = load_from_disk(str(self.config.test_dataset_path))
+        
+        # We DO NOT need transforms here because the data is already processed
+        # test_data.set_transform(...) # REMOVED
+
+        # Use the default collator which handles 'pixel_values' and 'label'
+        data_collator = DefaultDataCollator()
+
+        # Dummy trainer for running predictions
+        eval_args = TrainingArguments(
+            output_dir="./eval_output", # Temporary directory
+            per_device_eval_batch_size=self.config.batch_size,
+            report_to="none"
+        )
+        trainer = Trainer(
+            model=model,
+            args=eval_args,
+            data_collator=data_collator
+        )
+
+        # --- Run Predictions ---
+        logger.info("Running final evaluation on the test set...")
+        outputs = trainer.predict(test_data)
+        y_true = outputs.label_ids
+        y_pred = outputs.predictions.argmax(1)
+
+        # --- Calculate Metrics ---
+        scores = {
+            "accuracy": accuracy_score(y_true, y_pred),
+            "f1_score": f1_score(y_true, y_pred, average='macro'),
+            "precision": precision_score(y_true, y_pred, average='macro'),
+            "recall": recall_score(y_true, y_pred, average='macro')
+        }
+        logger.info(f"Test Set Metrics: {scores}")
+
+        # --- Save Metrics to a JSON file ---
+        metrics_path = Path(self.config.metrics_file_name) 
+        
+        # Now create the directory
+        metrics_path.parent.mkdir(parents=True, exist_ok=True)
+        
+        with open(metrics_path, 'w') as f:
+            json.dump(scores, f, indent=4)
+        logger.info(f"Metrics saved to {metrics_path}")
+        
+        # --- Log to MLflow ---
+        mlflow.set_tracking_uri(self.config.mlflow_uri)
+        mlflow.set_experiment("Pneumonia-ViT-Classification")
+
+        with mlflow.start_run():
+            logger.info("Logging parameters and metrics to MLflow...")
+            mlflow.log_params(self.config.all_params)
+            mlflow.log_metrics(scores)
+            
+            # --- THIS IS THE FINAL FIX ---
+            # Instead of logging the model object, log the directory where the
+            # trained model was already saved by the Trainer.
+            # `mlflow.log_artifact` is a simple upload and will not cause registry errors.
+            model_dir_path = str(self.config.path_of_model)
+            mlflow.log_artifact(model_dir_path, artifact_path="model")
+
+            logger.info("Successfully logged artifacts to MLflow.")

src/vitClassifier/components/model_training.py

@@ -0,0 +1,74 @@
+# src/vitClassifier/components/model_training.py
+
+import torch
+from datasets import load_from_disk
+from transformers import (ViTImageProcessor, ViTForImageClassification, TrainingArguments, Trainer, DefaultDataCollator)
+from vitClassifier.entity.config_entity import TrainingConfig
+from vitClassifier import logger
+import evaluate
+
+class ModelTraining:
+    def __init__(self, config: TrainingConfig):
+        self.config = config
+
+    def train(self):
+        # --- NEW: Explicitly define the device ---
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        logger.info(f"Using device: {device}")
+
+        # --- Load datasets (no change) ---
+        train_data = load_from_disk(str(self.config.train_dataset_path))
+        val_data = load_from_disk(str(self.config.val_dataset_path))
+        
+        id2label = {i: label for i, label in enumerate(train_data.features['label'].names)}
+        label2id = {label: i for i, label in id2label.items()}
+        
+        model = ViTForImageClassification.from_pretrained(
+            self.config.model_name, num_labels=len(id2label), id2label=id2label,
+            label2id=label2id, ignore_mismatched_sizes=True
+        )
+
+        # --- NEW: Move the model to the correct device ---
+        model.to(device)
+
+        # --- TrainingArguments (no change) ---
+        args = TrainingArguments(
+            output_dir=str(self.config.root_dir),
+            learning_rate=self.config.learning_rate,
+            per_device_train_batch_size=self.config.batch_size,
+            per_device_eval_batch_size=self.config.batch_size,
+            num_train_epochs=self.config.epochs,
+            weight_decay=self.config.weight_decay,
+            warmup_steps=self.config.warmup_steps,
+            save_strategy='epoch',
+            eval_strategy='epoch',
+            load_best_model_at_end=True,
+            metric_for_best_model="accuracy",
+            save_total_limit=1,
+            report_to="none"
+        )
+        
+        metric = evaluate.load("accuracy")
+        
+        def compute_metrics(eval_pred):
+            predictions, labels = eval_pred
+            predictions = predictions.argmax(axis=1)
+            return metric.compute(predictions=predictions, references=labels)
+        
+        data_collator = DefaultDataCollator()
+        processor = ViTImageProcessor.from_pretrained(self.config.model_name)
+
+        trainer = Trainer(
+            model, # The model is now already on the GPU
+            args,
+            train_dataset=train_data,
+            eval_dataset=val_data,
+            data_collator=data_collator,
+            compute_metrics=compute_metrics,
+            tokenizer=processor,
+        )
+
+        logger.info("Starting model fine-tuning with validation...")
+        trainer.train()
+        trainer.save_model(str(self.config.trained_model_path))
+        logger.info("Model fine-tuning complete and best model saved.")

src/vitClassifier/config/configuration.py

@@ -0,0 +1,76 @@
+# src/vitClassifier/config/configuration.py
+
+from vitClassifier.constants import CONFIG_FILE_PATH, PARAMS_FILE_PATH # <-- THIS IMPORT IS THE FIX
+from vitClassifier.utils.common import read_yaml, create_directories
+from vitClassifier.entity.config_entity import (DataIngestionConfig,
+                                                  DataTransformationConfig,
+                                                  TrainingConfig,
+                                                  EvaluationConfig)
+from pathlib import Path
+import os
+
+class ConfigurationManager:
+    def __init__(self, config_filepath=None, params_filepath=None):
+        
+        # If no path is provided when creating an instance, use the imported constants
+        if config_filepath is None:
+            config_filepath = CONFIG_FILE_PATH
+        if params_filepath is None:
+            params_filepath = PARAMS_FILE_PATH
+
+        self.config = read_yaml(config_filepath)
+        self.params = read_yaml(params_filepath)
+        create_directories([self.config.artifacts_root])
+
+    def get_data_ingestion_config(self) -> DataIngestionConfig:
+        config = self.config.data_ingestion
+        create_directories([config.root_dir])
+        return DataIngestionConfig(
+            root_dir=Path(config.root_dir),
+            source_kaggle_dataset_id=config.source_kaggle_dataset_id,
+            unzip_dir=Path(config.unzip_dir),
+            train_df_path=Path(config.train_df_path),
+            test_df_path=Path(config.test_df_path),
+            val_df_path=Path(config.val_df_path)
+        )
+
+    def get_data_transformation_config(self) -> DataTransformationConfig:
+        config = self.config.data_transformation
+        create_directories([config.root_dir])
+        return DataTransformationConfig(
+            root_dir=Path(config.root_dir),
+            train_data_path=Path(config.train_data_path),
+            test_data_path=Path(config.test_data_path),
+            val_data_path=Path(config.val_data_path),
+            train_dataset_path=Path(config.train_dataset_path),
+            test_dataset_path=Path(config.test_dataset_path),
+            val_dataset_path=Path(config.val_dataset_path)
+        )
+    
+    def get_training_config(self) -> TrainingConfig:
+        training = self.config.model_training
+        params = self.params
+        create_directories([Path(training.root_dir)])
+        return TrainingConfig(
+            root_dir=Path(training.root_dir),
+            trained_model_path=Path(training.trained_model_path),
+            model_name=training.model_name,
+            train_dataset_path=Path(training.train_dataset_path),
+            val_dataset_path=Path(training.val_dataset_path),
+            learning_rate=params.LEARNING_RATE,
+            batch_size=params.BATCH_SIZE,
+            epochs=params.EPOCHS,
+            weight_decay=params.WEIGHT_DECAY,
+            warmup_steps=params.WARMUP_STEPS,
+        )
+
+    def get_evaluation_config(self) -> EvaluationConfig:
+        eval_config = self.config.model_evaluation
+        return EvaluationConfig(
+            path_of_model=Path(eval_config.model_path),
+            test_dataset_path=Path(eval_config.test_dataset_path),
+            mlflow_uri=eval_config.mlflow_uri,
+            all_params=self.params,
+            batch_size=self.params.BATCH_SIZE,
+            metrics_file_name=Path(eval_config.metrics_file_name) # <--- MAKE SURE THIS LINE EXISTS
+        )

src/vitClassifier/constants/__init__.py

@@ -0,0 +1,4 @@
+from pathlib import Path
+
+CONFIG_FILE_PATH = Path("config/config.yaml")
+PARAMS_FILE_PATH = Path("params.yaml")

src/vitClassifier/entity/config_entity.py

@@ -0,0 +1,43 @@
+from dataclasses import dataclass
+from pathlib import Path
+
+@dataclass(frozen=True)
+class DataIngestionConfig:
+    root_dir: Path
+    source_kaggle_dataset_id: str
+    unzip_dir: Path
+    train_df_path: Path # New
+    test_df_path: Path  # New
+    val_df_path: Path   # New
+
+@dataclass(frozen=True)
+class DataTransformationConfig:
+    root_dir: Path
+    train_data_path: Path # New
+    test_data_path: Path  # New
+    val_data_path: Path   # New
+    train_dataset_path: Path
+    test_dataset_path: Path
+    val_dataset_path: Path # New
+
+@dataclass(frozen=True)
+class TrainingConfig:
+    root_dir: Path
+    trained_model_path: Path
+    model_name: str
+    train_dataset_path: Path # New
+    val_dataset_path: Path   # New
+    learning_rate: float
+    batch_size: int
+    epochs: int
+    weight_decay: float
+    warmup_steps: int
+
+@dataclass(frozen=True)
+class EvaluationConfig:
+    path_of_model: Path
+    test_dataset_path: Path
+    mlflow_uri: str
+    all_params: dict
+    batch_size: int
+    metrics_file_name: Path

src/vitClassifier/pipeline/stage_01_data_ingestion.py

@@ -0,0 +1,27 @@
+# src/vitClassifier/pipeline/stage_01_data_ingestion.py
+
+from vitClassifier.config.configuration import ConfigurationManager
+from vitClassifier.components.data_ingestion import DataIngestion
+from vitClassifier import logger
+
+STAGE_NAME = "Data Ingestion stage"
+
+class DataIngestionTrainingPipeline:
+    def __init__(self):
+        pass
+    def main(self):
+        config = ConfigurationManager()
+        data_ingestion_config = config.get_data_ingestion_config()
+        data_ingestion = DataIngestion(config=data_ingestion_config)
+        data_ingestion.ingest_data()
+
+# <<< ADD THIS BLOCK TO MAKE THE SCRIPT RUNNABLE >>>
+if __name__ == '__main__':
+    try:
+        logger.info(f">>>>>> stage {STAGE_NAME} started <<<<<<")
+        obj = DataIngestionTrainingPipeline()
+        obj.main()
+        logger.info(f">>>>>> stage {STAGE_NAME} completed <<<<<<\n\nx==========x")
+    except Exception as e:
+        logger.exception(e)
+        raise e

src/vitClassifier/pipeline/stage_02_data_transformation.py

@@ -0,0 +1,32 @@
+from vitClassifier.config.configuration import ConfigurationManager
+from vitClassifier.components.data_transformation import DataTransformation
+from vitClassifier import logger
+
+STAGE_NAME = "Data Transformation stage"
+
+class DataTransformationTrainingPipeline:
+    def __init__(self):
+        pass
+    def main(self):
+        config_manager = ConfigurationManager()
+        data_transformation_config = config_manager.get_data_transformation_config()
+        params = config_manager.params
+        # Get model_name from the training config section
+        model_name = config_manager.config.model_training.model_name
+        
+        data_transformation = DataTransformation(
+            config=data_transformation_config,
+            random_state=params.RANDOM_STATE,
+            model_name=model_name # Pass the model name
+        )
+        data_transformation.transform_data()
+
+if __name__ == '__main__':
+    try:
+        logger.info(f">>>>>> stage {STAGE_NAME} started <<<<<<")
+        obj = DataTransformationTrainingPipeline()
+        obj.main()
+        logger.info(f">>>>>> stage {STAGE_NAME} completed <<<<<<\n\nx==========x")
+    except Exception as e:
+        logger.exception(e)
+        raise e

src/vitClassifier/pipeline/stage_03_model_training.py

@@ -0,0 +1,24 @@
+from vitClassifier.config.configuration import ConfigurationManager
+from vitClassifier.components.model_training import ModelTraining
+from vitClassifier import logger
+
+STAGE_NAME = "Model Training stage"
+
+class ModelTrainingPipeline:
+    def __init__(self):
+        pass
+    def main(self):
+        config = ConfigurationManager()
+        training_config = config.get_training_config()
+        model_training = ModelTraining(config=training_config)
+        model_training.train()
+
+if __name__ == '__main__':
+    try:
+        logger.info(f">>>>>> stage {STAGE_NAME} started <<<<<<")
+        obj = ModelTrainingPipeline()
+        obj.main()
+        logger.info(f">>>>>> stage {STAGE_NAME} completed <<<<<<\n\nx==========x")
+    except Exception as e:
+        logger.exception(e)
+        raise e

src/vitClassifier/pipeline/stage_04_model_evaluation.py

@@ -0,0 +1,26 @@
+from vitClassifier.config.configuration import ConfigurationManager
+from vitClassifier.components.model_evaluation import ModelEvaluation
+from vitClassifier import logger
+from dotenv import load_dotenv
+load_dotenv()
+
+STAGE_NAME = "Model Evaluation stage"
+
+class ModelEvaluationPipeline:
+    def __init__(self):
+        pass
+    def main(self):
+        config = ConfigurationManager()
+        eval_config = config.get_evaluation_config()
+        evaluation = ModelEvaluation(config=eval_config)
+        evaluation.evaluate()
+
+if __name__ == '__main__':
+    try:
+        logger.info(f">>>>>> stage {STAGE_NAME} started <<<<<<")
+        obj = ModelEvaluationPipeline()
+        obj.main()
+        logger.info(f">>>see stage {STAGE_NAME} completed <<<<<<\n\nx==========x")
+    except Exception as e:
+        logger.exception(e)
+        raise e

src/vitClassifier/utils/common.py

@@ -0,0 +1,28 @@
+import os
+import yaml
+import json
+from box import ConfigBox
+from box.exceptions import BoxValueError
+from ensure import ensure_annotations
+from pathlib import Path
+from typing import Any
+from vitClassifier import logger
+
+@ensure_annotations
+def read_yaml(path_to_yaml: Path) -> ConfigBox:
+    try:
+        with open(path_to_yaml) as yaml_file:
+            content = yaml.safe_load(yaml_file)
+            logger.info(f"yaml file: {path_to_yaml} loaded successfully")
+            return ConfigBox(content)
+    except BoxValueError:
+        raise ValueError("yaml file is empty")
+    except Exception as e:
+        raise e
+
+@ensure_annotations
+def create_directories(path_to_directories: list, verbose=True):
+    for path in path_to_directories:
+        os.makedirs(path, exist_ok=True)
+        if verbose:
+            logger.info(f"created directory at: {path}")

template.py

@@ -0,0 +1,40 @@
+import os
+from pathlib import Path
+import logging
+
+logging.basicConfig(level=logging.INFO, format='[%(asctime)s]: %(message)s:')
+
+project_name = "vitClassifier"
+
+list_of_files = [
+    ".github/workflows/.gitkeep",
+    f"src/{project_name}/__init__.py",
+    f"src/{project_name}/components/__init__.py",
+    f"src/{project_name}/utils/__init__.py",
+    f"src/{project_name}/config/__init__.py",
+    f"src/{project_name}/config/configuration.py",
+    f"src/{project_name}/pipeline/__init__.py",
+    f"src/{project_name}/entity/__init__.py",
+    f"src/{project_name}/constants/__init__.py",
+    "config/config.yaml",
+    "dvc.yaml",
+    "params.yaml",
+    "requirements.txt",
+    "setup.py",
+    "research/trials.ipynb"
+]
+
+for filepath in list_of_files:
+    filepath = Path(filepath)
+    filedir, filename = os.path.split(filepath)
+
+    if filedir != "":
+        os.makedirs(filedir, exist_ok=True)
+        logging.info(f"Creating directory; {filedir} for the file: {filename}")
+
+    if (not os.path.exists(filepath)) or (os.path.getsize(filepath) == 0):
+        with open(filepath, "w") as f:
+            pass
+            logging.info(f"Creating empty file: {filepath}")
+    else:
+        logging.info(f"{filename} is already exists")