Initial Commit

.gitattributes

@@ -0,0 +1 @@
+weights/efficientvit.pth filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,182 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.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/
+*.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
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# 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
+.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/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Data
+Data/
+
+EDA/
+
+# Model Outputs
+model_outputs/

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

.python-version

@@ -0,0 +1 @@
+3.12.9

DockerFile

@@ -0,0 +1,28 @@
+FROM python:3.12-slim
+
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    git \
+    libgl1-mesa-glx \
+    libglib2.0-0 \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy requirements first to leverage Docker cache
+COPY requirements.txt .
+
+# Install Python dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy project files
+COPY . .
+
+# Make the gradio inference script executable
+RUN chmod +x gradio_inference.py
+
+# Expose port for Gradio
+EXPOSE 7860
+
+# Set the entrypoint command
+CMD ["python", "gradio_inference.py"]

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -0,0 +1,123 @@
+# Eye Disease Detection 
+
+This repository contains a Gradio web application for eye disease detection using deep learning models. The application allows users to upload fundus images and get predictions for common eye conditions.
+
+## Features
+
+- **Easy-to-use web interface** for eye disease detection
+- Support for **multiple model architectures** (MobileNetV4, LeViT, EfficientViT, GENet, RegNetX)
+- **Custom model loading** from saved model checkpoints
+- **Visualization** of prediction probabilities
+- **Dockerized deployment** option
+
+## Supported Eye Conditions
+
+The system can detect the following eye conditions:
+- Central Serous Chorioretinopathy
+- Diabetic Retinopathy
+- Disc Edema
+- Glaucoma
+- Healthy (normal eye)
+- Macular Scar
+- Myopia
+- Retinal Detachment
+- Retinitis Pigmentosa
+
+## Installation
+
+### Prerequisites
+
+- Python 3.12+
+- PyTorch 2.7.0+
+- CUDA-compatible GPU (optional, but recommended for faster inference)
+
+### Option 1: Local Installation
+
+1. Clone this repository:
+   ```bash
+   git clone https://github.com/GilbertKrantz/eye-disease-detection.git
+   cd eye-disease-detection
+   ```
+
+2. Install the required packages:
+   ```bash
+   pip install -r requirements.txt
+   ```
+
+3. Run the application:
+   ```bash
+   python gradio_inference.py
+   ```
+
+4. Open your browser and go to http://localhost:7860
+
+### Option 2: Docker Installation
+
+1. Build the Docker image:
+   ```bash
+   docker build -t eye-disease-detection .
+   ```
+
+2. Run the container:
+   ```bash
+   docker run -p 7860:7860 eye-disease-detection
+   ```
+
+3. Open your browser and go to http://localhost:7860
+
+## Usage
+
+1. Upload a fundus image of the eye
+2. (Optional) Specify the path to your trained model file (.pth)
+3. Select the model architecture (MobileNetV4, LeViT, EfficientViT, GENet, RegNetX)
+4. Click "Analyze Image" to get the prediction
+5. View the results and probability distribution
+
+## Model Training
+
+This repository focuses on inference. For training your own models, refer to the main training script and follow these steps:
+
+1. Prepare your dataset in the required directory structure
+2. Train a model using the main.py script:
+   ```bash
+   python main.py --train-dir "/path/to/training/data" --eval-dir "/path/to/eval/data" --model mobilenetv4 --epochs 20 --save-model "my_model.pth"
+   ```
+3. Use the saved model with the inference application
+
+## Project Structure
+
+```
+.
+├── gradio_inference.py     # Main Gradio application
+├── requirements.txt        # Python dependencies
+├── Dockerfile              # Docker configuration
+├── README.md               # This documentation
+├── utils/                  # Utility modules
+│   ├── ModelCreator.py     # Model architecture definitions
+│   ├── Evaluator.py        # Model evaluation utilities
+│   ├── DatasetHandler.py   # Dataset handling utilities
+│   ├── Trainer.py          # Model training utilities
+│   └── Callback.py         # Training callbacks
+└── main.py                 # Main training script
+```
+
+## Performance
+
+The performance of the models depends on the quality of training data and the specific architecture used. In general, these models can achieve accuracy rates of 85-95% on standard eye disease datasets.
+
+## Customization
+
+You can customize the application in several ways:
+- Add example images in the Gradio interface
+- Extend the list of supported classes by modifying the CLASSES variable in gradio_inference.py
+- Add support for additional model architectures in ModelCreator.py
+
+## License
+
+This project is licensed under the Apache License 2.0 - see the LICENSE file for details.
+
+## Acknowledgments
+
+- The models are built using PyTorch and the TIMM library
+- The web interface is built using Gradio
+- Special thanks to the open-source community for making this project possible

gradio-inference.py

@@ -0,0 +1,221 @@
+#!/usr/bin/env python
+# Eye Disease Detection - Gradio Inference App
+# Date: May 11, 2025
+
+import os
+import sys
+import torch
+import numpy as np
+import gradio as gr
+from PIL import Image
+from torchvision import transforms
+import logging
+
+logging.basicConfig(level=logging.INFO)
+
+# Import custom modules
+sys.path.append("./utils")
+from ModelCreator import EyeDetectionModels
+
+# Set device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+# Define class names (make sure these match your model's classes)
+CLASSES = [
+    "Central Serous Chorioretinopathy",
+    "Diabetic Retinopathy",
+    "Disc Edema",
+    "Glaucoma",
+    "Healthy",
+    "Macular Scar",
+    "Myopia",
+    "Retinal Detachment",
+    "Retinitis Pigmentosa",
+]
+
+
+def get_transform():
+    """Get the standard transformation pipeline for inference."""
+    return transforms.Compose(
+        [
+            transforms.Resize(256),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]
+    )
+
+
+def load_model(model_path, model_type="efficientvit"):
+    """
+    Load a pretrained model for inference.
+
+    Args:
+        model_path: Path to the saved model state dict
+        model_type: Type of model to load (mobilenetv4, levit, efficientvit, gernet, regnetx)
+
+    Returns:
+        Loaded model ready for inference
+    """
+    # Initialize model creator
+    logging.info("Initializing model creator...")
+    model_creator = EyeDetectionModels(
+        num_classes=len(CLASSES), freeze_layers=False  # Not relevant for inference
+    )
+
+    # Check if model type exists
+    if model_type not in model_creator.models:
+        raise ValueError(
+            f"Model type '{model_type}' not found. Available models: {list(model_creator.models.keys())}"
+        )
+
+    # Create model of specified type
+    logging.info(f"Creating model of type: {model_type}")
+    model = model_creator.models[model_type]()
+
+    # Load state dict if provided
+    if model_path and not os.path.exists(model_path):
+        raise FileNotFoundError(f"Model path '{model_path}' does not exist.")
+    elif model_path is None:
+        #  Use default model path if it exists
+        if os.path.exists(f"./weights/{model_type}.pth"):
+            model_path = f"./weights/{model_type}.pth"
+        else:
+            model_path = None
+            logging.warning(
+                f"Default model path '{model_path}' not found. Using untrained model."
+            )
+    # Set model to evaluation mode
+    model.eval()
+    return model
+
+
+def predict_image(image, model_path, model_type):
+    """
+    Predict eye disease from an uploaded image.
+
+    Args:
+        image: Input image from Gradio
+        model_path: Path to the model state dict
+        model_type: Type of model architecture
+
+    Returns:
+        Dictionary of class probabilities
+    """
+    try:
+
+        logging.info("Starting prediction...")
+        # Load model
+        model = load_model(model_path, model_type)
+
+        # Preprocess image
+        logging.info("Preprocessing image...")
+        if image is None:
+            logging.warning("No image provided.")
+            return {cls: 0.0 for cls in CLASSES}
+        transform = get_transform()
+        if image is None:
+            return {cls: 0.0 for cls in CLASSES}
+
+        # Convert numpy array to PIL Image
+        img = Image.fromarray(image).convert("RGB")
+        img_tensor = transform(img).unsqueeze(0).to(device)
+        logging.info("Image preprocessed successfully.")
+
+        # Make prediction
+        with torch.no_grad():
+            outputs = model(img_tensor)
+            probabilities = torch.nn.functional.softmax(outputs, dim=1)[0].cpu().numpy()
+
+        # Return probabilities for each class
+        return {cls: float(prob) for cls, prob in zip(CLASSES, probabilities)}
+
+    except Exception as e:
+        import traceback
+
+        traceback.print_exc()
+        return {cls: 0.0 for cls in CLASSES}
+
+
+def main():
+    """Main function to run the Gradio interface."""
+    # Define available models
+    model_types = ["mobilenetv4", "levit", "efficientvit", "gernet", "regnetx"]
+
+    # Create the Gradio interface
+    with gr.Blocks(title="Eye Disease Detection") as demo:
+        gr.Markdown("# Eye Disease Detection System")
+        gr.Markdown(
+            """This application uses deep learning to detect eye diseases from fundus images.
+                    Currently supports detection of: Cataract, Diabetic Retinopathy, Glaucoma, and Normal eyes."""
+        )
+
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(label="Upload Fundus Image", type="numpy")
+                model_path = gr.Textbox(
+                    label="Model Path (leave empty to use default)",
+                    placeholder="Path to model .pth file",
+                    value="",
+                )
+                model_type = gr.Dropdown(
+                    label="Model Architecture", choices=model_types, value="mobilenetv4"
+                )
+                submit_btn = gr.Button("Analyze Image", variant="primary")
+
+            with gr.Column():
+                output_chart = gr.Label(label="Prediction")
+
+        # Process the image when the button is clicked
+        submit_btn.click(
+            fn=predict_image,
+            inputs=[input_image, model_path, model_type],
+            outputs=output_chart,
+        )
+
+        # Examples section
+        gr.Markdown("### Examples (Please add your own example images)")
+        gr.Examples(
+            examples=[],  # Add example paths here
+            inputs=input_image,
+            outputs=[output_chart],
+            fn=predict_image,
+            cache_examples=True,
+        )
+
+        # Usage instructions
+        with gr.Accordion("Usage Instructions", open=False):
+            gr.Markdown(
+                """
+            ## How to use this application:
+            
+            1. **Upload an image**: Click the upload button to select a fundus image from your computer
+            2. **Specify model** (Optional): 
+                - Enter the path to your trained model file (.pth)
+                - Select the model architecture that was used for training
+            3. **Analyze**: Click the "Analyze Image" button to get results
+            4. **Interpret results**: The system will show the detected condition and probability distribution
+            
+            ## Model Information:
+            
+            This system supports multiple model architectures:
+            - **MobileNetV4**: Lightweight and efficient model
+            - **LeViT**: Vision Transformer designed for efficiency
+            - **EfficientViT**: Hybrid CNN-Transformer architecture
+            - **GENet**: General and Efficient Network
+            - **RegNetX**: Systematically designed CNN architecture
+            
+            For best results, ensure you're using a high-quality fundus image and the correct model type.
+            """
+            )
+
+    # Launch the app
+    demo.launch(
+        share=True,
+        pwa=True,
+    )
+
+
+if __name__ == "__main__":
+    main()

main.py

@@ -0,0 +1,342 @@
+#!/usr/bin/env python
+# Eye Disease Detection - Main Application
+# Date: May 11, 2025
+
+import os
+import sys
+import argparse
+import random
+import numpy as np
+import torch
+import torch.nn as nn
+import matplotlib.pyplot as plt
+from torchvision import transforms, datasets
+from torch.utils.data import DataLoader, random_split
+
+# Import custom modules
+sys.path.append("./utils")
+from ModelCreator import EyeDetectionModels
+from DatasetHandler import FilteredImageDataset
+from Evaluator import ClassificationEvaluator
+from Comparator import compare_models
+from Trainer import model_train
+
+
+# Set random seeds for reproducibility
+def set_seed(seed=42):
+    """Set seeds for reproducibility."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+
+
+def get_transform():
+    """
+    Get standard data transform for both training and validation/testing.
+
+    Returns:
+        transform: Standard transform for all datasets
+    """
+    # Standard transform as specified
+    transform = transforms.Compose(
+        [
+            transforms.Resize(256),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]
+    )
+
+    return transform
+
+
+def load_data(args):
+    """
+    Load and prepare datasets from separate directories for training and evaluation.
+
+    Args:
+        args: Command line arguments
+
+    Returns:
+        train_loader: DataLoader for training
+        val_loader: DataLoader for validation
+        test_loader: DataLoader for testing
+        dataset_ref: Reference to the evaluation dataset for class information
+    """
+    print(f"Loading training dataset from: {args.train_dir}")
+    print(f"Loading evaluation dataset from: {args.eval_dir}")
+
+    # Get standard transform
+    transform = get_transform()
+
+    # Load training dataset
+    train_dataset = datasets.ImageFolder(args.train_dir, transform=transform)
+    print(f"Training dataset classes: {train_dataset.classes}")
+    print(f"Training dataset size: {len(train_dataset)}")
+
+    # Load evaluation dataset
+    eval_dataset = datasets.ImageFolder(args.eval_dir, transform=transform)
+    print(f"Evaluation dataset classes: {eval_dataset.classes}")
+
+    # Apply class filtering if requested
+    excluded_classes = args.exclude_classes.split(",") if args.exclude_classes else None
+    if excluded_classes and any(excluded_classes):
+        train_dataset = FilteredImageDataset(train_dataset, excluded_classes)
+        eval_dataset = FilteredImageDataset(eval_dataset, excluded_classes)
+        print(f"After filtering - Classes: {eval_dataset.classes}")
+
+    print(f"After filtering - Train size: {len(train_dataset)}")
+    print(f"After filtering - Eval size: {len(eval_dataset)}")
+
+    # Split evaluation dataset into validation and test sets
+    val_size = int(
+        len(eval_dataset) * (args.val_split / (args.val_split + args.test_split))
+    )
+    test_size = len(eval_dataset) - val_size
+
+    val_dataset, test_dataset = random_split(eval_dataset, [val_size, test_size])
+
+    print(
+        f"Split sizes - Train: {len(train_dataset)}, "
+        f"Validation: {len(val_dataset)}, Test: {len(test_dataset)}"
+    )
+
+    # Create data loaders
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+        pin_memory=True,
+    )
+
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=args.num_workers,
+        pin_memory=True,
+    )
+
+    test_loader = DataLoader(
+        test_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=args.num_workers,
+        pin_memory=True,
+    )
+
+    # Use eval_dataset as the reference for class information
+    return train_loader, val_loader, test_loader, eval_dataset
+
+
+def train_single_model(args, train_loader, val_loader, test_loader, dataset):
+    """Train a single model specified by the arguments."""
+
+    print(f"Creating {args.model} model...")
+
+    # Initialize model creator
+    model_creator = EyeDetectionModels(
+        num_classes=len(dataset.classes), freeze_layers=(not args.unfreeze_all)
+    )
+
+    # Get model
+    if args.model in model_creator.models:
+        model = model_creator.models[args.model]()
+    else:
+        available_models = list(model_creator.models.keys())
+        print(
+            f"Error: Model '{args.model}' not found. Available models: {available_models}"
+        )
+        sys.exit(1)
+
+    # Train and evaluate model
+    results = model_train(model, train_loader, val_loader, dataset, epochs=args.epochs)
+
+    # Test the model
+    if results["accuracy"] is not None:
+        print("\nEvaluating on test set...")
+        evaluator = ClassificationEvaluator(class_names=dataset.classes)
+        test_results = evaluator.evaluate_model(model, test_loader)
+        print(f"Test accuracy: {test_results['accuracy']:.4f}")
+
+        # Save model if requested
+        if args.save_model:
+            save_path = args.save_model
+            try:
+                torch.save(model.state_dict(), save_path)
+                print(f"Model saved to {save_path}")
+            except Exception as e:
+                print(f"Error saving model: {e}")
+    else:
+        print("Training failed. Cannot evaluate on test set.")
+
+
+def compare_multiple_models(args, train_loader, val_loader, test_loader, dataset):
+    """Compare multiple models."""
+
+    print("Preparing to compare multiple models...")
+
+    # Initialize model creator
+    model_creator = EyeDetectionModels(
+        num_classes=len(dataset.classes), freeze_layers=(not args.unfreeze_all)
+    )
+
+    # Get list of models to compare
+    model_names = args.compare_models.split(",")
+    models = []
+    names = []
+
+    for model_name in model_names:
+        model_name = model_name.strip()
+        if model_name in model_creator.models:
+            print(f"Adding {model_name} to comparison...")
+            models.append(model_creator.models[model_name]())
+            names.append(model_name)
+        else:
+            print(f"Warning: Model '{model_name}' not found, skipping.")
+
+    if not models:
+        print("No valid models to compare. Exiting.")
+        return
+
+    # Run comparison
+    compare_models(
+        models,
+        train_loader,
+        val_loader,
+        test_loader,
+        dataset,
+        epochs=args.epochs,
+        names=names,
+    )
+
+
+def main():
+    """Main function to run the eye disease detection application."""
+
+    # Set up argument parser with example usage
+    parser = argparse.ArgumentParser(
+        description="Eye Disease Detection using Deep Learning",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+        epilog="""
+Examples:
+  # Train a single model
+  python main.py --train-dir "/path/to/augmented_dataset" --eval-dir "/path/to/original_dataset" --model mobilenetv4 --epochs 20 --save-model best_model.pth
+  
+  # Compare multiple models
+  python main.py --train-dir "/path/to/augmented_dataset" --eval-dir "/path/to/original_dataset" --compare-models mobilenetv4,levit,efficientvit --epochs 15
+""",
+    )
+
+    # Dataset and data loading arguments
+    data_group = parser.add_argument_group("Data Options")
+    data_group.add_argument(
+        "--train-dir",
+        type=str,
+        required=True,
+        help="Path to the training dataset directory (Augmented Dataset)",
+    )
+    data_group.add_argument(
+        "--eval-dir",
+        type=str,
+        required=True,
+        help="Path to the evaluation dataset directory (Original Dataset)",
+    )
+    data_group.add_argument(
+        "--batch-size",
+        type=int,
+        default=32,
+        help="Batch size for training and evaluation",
+    )
+    data_group.add_argument(
+        "--val-split",
+        type=float,
+        default=0.5,
+        help="Validation split ratio within evaluation set",
+    )
+    data_group.add_argument(
+        "--test-split",
+        type=float,
+        default=0.5,
+        help="Test split ratio within evaluation set",
+    )
+    data_group.add_argument(
+        "--num-workers",
+        type=int,
+        default=4,
+        help="Number of worker processes for data loading",
+    )
+    data_group.add_argument(
+        "--exclude-classes",
+        type=str,
+        default=None,
+        help="Comma-separated list of class names to exclude",
+    )
+
+    # Model arguments
+    model_group = parser.add_argument_group("Model Options")
+    model_group.add_argument(
+        "--model",
+        type=str,
+        default="mobilenetv4",
+        help="Model architecture to use. Options: mobilenetv4, levit, efficientvit, gernet, regnetx",
+    )
+    model_group.add_argument(
+        "--unfreeze-all", action="store_true", help="Unfreeze all layers for training"
+    )
+    model_group.add_argument(
+        "--compare-models",
+        type=str,
+        default=None,
+        help="Comma-separated list of models to compare",
+    )
+
+    # Training arguments
+    train_group = parser.add_argument_group("Training Options")
+    train_group.add_argument(
+        "--epochs", type=int, default=20, help="Number of training epochs"
+    )
+    train_group.add_argument(
+        "--seed", type=int, default=42, help="Random seed for reproducibility"
+    )
+    train_group.add_argument(
+        "--save-model", type=str, default=None, help="Path to save the trained model"
+    )
+
+    # Parse arguments
+    args = parser.parse_args()
+
+    # Set random seed for reproducibility
+    set_seed(args.seed)
+
+    # Display GPU information
+    if torch.cuda.is_available():
+        device_count = torch.cuda.device_count()
+        print(f"Using {device_count} GPU{'s' if device_count > 1 else ''}")
+        for i in range(device_count):
+            print(f"  Device {i}: {torch.cuda.get_device_name(i)}")
+    else:
+        print("No GPU available, using CPU")
+
+    # Load data
+    train_loader, val_loader, test_loader, dataset = load_data(args)
+
+    # Check if comparing multiple models
+    if args.compare_models:
+        compare_multiple_models(args, train_loader, val_loader, test_loader, dataset)
+    else:
+        train_single_model(args, train_loader, val_loader, test_loader, dataset)
+
+
+if __name__ == "__main__":
+    # Example usage for direct execution:
+    # python main.py --train-dir "/kaggle/input/eye-disease-image-dataset/Augmented Dataset/Augmented Dataset" \
+    #                --eval-dir "/kaggle/input/eye-disease-image-dataset/Original Dataset/Original Dataset" \
+    #                --model mobilenetv4 --epochs 10
+    main()

pyproject.toml

@@ -0,0 +1,36 @@
+[project]
+name = "eyediseasedetection"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.12.9"
+dependencies = [
+  "gradio>=5.29.0",
+  "matplotlib>=3.10.3",
+  "pandas>=2.2.3",
+  "scikit-learn>=1.6.1",
+  "seaborn>=0.13.2",
+  "timm>=1.0.15",
+  "torch>=2.7.0",
+  "torchaudio>=2.7.0",
+  "torchvision>=0.22.0",
+  "tqdm>=4.67.1",
+]
+
+[tool.uv.sources]
+torch = [
+  { index = "pytorch-cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+]
+torchvision = [
+  { index = "pytorch-cu128", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+]
+
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
+[[tool.uv.index]]
+name = "pytorch-cu128"
+url = "https://download.pytorch.org/whl/cu128"
+explicit = true

requirements.txt

@@ -0,0 +1,11 @@
+torch>=2.7.0
+torchvision>=0.22.0
+torchaudio>=2.7.0
+numpy>=1.26.0
+Pillow>=10.0.0
+gradio>=4.11.0
+matplotlib>=3.8.0
+seaborn>=0.13.0
+scikit-learn>=1.4.0
+tqdm>=4.66.0
+timm>=1.0.0

training.ipynb

@@ -0,0 +1 @@
+{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":10951558,"sourceType":"datasetVersion","datasetId":6812365}],"dockerImageVersionId":31011,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport time\nimport random\nimport copy\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nimport torchvision.transforms as transforms\nfrom torchvision import transforms, datasets\nimport torchvision.models as models\nimport timm\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom PIL import Image\nfrom torch.utils.data import DataLoader, random_split, Subset, Dataset\nfrom sklearn.metrics import (\n    accuracy_score, confusion_matrix, classification_report, roc_curve, auc, precision_recall_curve, average_precision_score\n)\nfrom sklearn.preprocessing import label_binarize\nfrom tqdm import tqdm\nimport gc\n\n\n# Set device\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Using device: {device}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-07T15:22:11.680667Z","iopub.execute_input":"2025-05-07T15:22:11.680991Z","iopub.status.idle":"2025-05-07T15:22:16.455808Z","shell.execute_reply.started":"2025-05-07T15:22:11.680969Z","shell.execute_reply":"2025-05-07T15:22:16.455137Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Modified FilteredImageDataset class with Pterygium filtering\nclass FilteredImageDataset(Dataset):\n    def __init__(self, dataset, excluded_classes=None):\n        \"\"\"\n        Create a filtered dataset that excludes specific classes.\n        \n        Args:\n            dataset: Original dataset (ImageFolder or similar)\n            excluded_classes: List of class names to exclude (e.g., [\"Pterygium\"])\n        \"\"\"\n        self.dataset = dataset\n        self.excluded_classes = excluded_classes or []\n        \n        # Get original class information\n        self.orig_classes = dataset.classes\n        self.orig_class_to_idx = dataset.class_to_idx\n        \n        # Create indices of samples to keep (excluding specified classes)\n        self.indices = []\n        for idx, (_, target) in enumerate(dataset.samples):\n            class_name = self.orig_classes[target]\n            if class_name not in self.excluded_classes:\n                self.indices.append(idx)\n        \n        # Create new class mapping without excluded classes\n        remaining_classes = [c for c in self.orig_classes if c not in self.excluded_classes]\n        self.classes = remaining_classes\n        self.class_to_idx = {cls: idx for idx, cls in enumerate(remaining_classes)}\n        self.idx_to_class = {v: k for k, v in self.class_to_idx.items()}\n        \n        # Create a mapping from old indices to new indices\n        self.target_mapping = {}\n        for old_class, old_idx in self.orig_class_to_idx.items():\n            if old_class in self.class_to_idx:\n                self.target_mapping[old_idx] = self.class_to_idx[old_class]\n        \n        print(f\"Filtered out classes: {self.excluded_classes}\")\n        print(f\"Remaining classes: {self.classes}\")\n        print(f\"Original dataset size: {len(dataset)}, Filtered dataset size: {len(self.indices)}\")\n\n    def __getitem__(self, index):\n        \"\"\"Get item from the filtered dataset with remapped class labels.\"\"\"\n        orig_idx = self.indices[index]\n        img, old_target = self.dataset[orig_idx]\n        \n        # Remap target to new class index\n        new_target = self.target_mapping[old_target]\n        \n        return img, new_target\n\n    def __len__(self):\n        \"\"\"Return the number of samples in the filtered dataset.\"\"\"\n        return len(self.indices)\n    \n    # Allow transform to be updated\n    def set_transform(self, transform):\n        \"\"\"Update the transform for the dataset.\"\"\"\n        self.dataset.transform = transform","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-07T15:22:16.457202Z","iopub.execute_input":"2025-05-07T15:22:16.457579Z","iopub.status.idle":"2025-05-07T15:22:16.465134Z","shell.execute_reply.started":"2025-05-07T15:22:16.457560Z","shell.execute_reply":"2025-05-07T15:22:16.464502Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Model Definition","metadata":{}},{"cell_type":"code","source":"# Early stopping class\nclass EarlyStopping:\n    def __init__(self, patience=5, delta=0):\n        self.patience = patience\n        self.delta = delta\n        self.counter = 0\n        self.best_score = None\n        self.early_stop = False\n        \n    def __call__(self, val_loss):\n        score = -val_loss\n        \n        if self.best_score is None:\n            self.best_score = score\n        elif score < self.best_score + self.delta:\n            self.counter += 1\n            if self.counter >= self.patience:\n                self.early_stop = True\n        else:\n            self.best_score = score\n            self.counter = 0\n\n# Model architecture functions\ndef _get_feature_blocks(model):\n    \"\"\"\n    Utility: locate the main feature blocks container in a timm model.\n    Returns a list-like module of blocks.\n    \"\"\"\n    for attr in ('features', 'blocks', 'layers', 'stem'):  # common container names\n        if hasattr(model, attr):\n            return getattr(model, attr)\n    # fallback: collect all children except classifier/head\n    return list(model.children())[:-1]\n\ndef _freeze_except_last_n(blocks, n):\n    total = len(blocks)\n    for idx, block in enumerate(blocks):\n        requires = (idx >= total - n)\n        for p in block.parameters():\n            p.requires_grad = requires\n\ndef get_model_mobilenetv4(num_classes, freeze_layers=True, device='cuda'):\n    model = timm.create_model('mobilenetv4_conv_medium.e500_r256_in1k', pretrained=True)\n    if freeze_layers:\n        blocks = _get_feature_blocks(model)\n        _freeze_except_last_n(blocks, 2)\n    # replace classifier\n    in_features = model.classifier.in_features\n    model.classifier = nn.Sequential(\n        nn.Linear(in_features, 512),\n        nn.ReLU(inplace=True),\n        nn.Dropout(0.4),\n        nn.Linear(512, num_classes)\n    )\n    return model.to(device)\n\ndef get_model_levit(num_classes, freeze_layers=True, device='cuda'):\n    model = timm.create_model('levit_128s.fb_dist_in1k', pretrained=True)\n    if freeze_layers:\n        blocks = _get_feature_blocks(model)\n        _freeze_except_last_n(blocks, 2)\n    # Attempt to extract in_features from model.head or classifier\n    head = getattr(model, 'head_dist', None) or getattr(model, 'classifier', None)\n    linear = getattr(head, 'linear')\n    in_features = 384\n    model.head = nn.Sequential(\n        nn.Linear(in_features, 512),\n        nn.ReLU(inplace=True),\n        nn.Dropout(0.4),\n        nn.Linear(512, num_classes)\n    )\n    model.head_dist = nn.Sequential(\n        nn.Linear(in_features, 512),\n        nn.ReLU(inplace=True),\n        nn.Dropout(0.4),\n        nn.Linear(512, num_classes)\n    )\n    return model.to(device)\n\ndef get_model_efficientvit(num_classes, freeze_layers=True, device='cuda'):\n    model = timm.create_model('efficientvit_m1.r224_in1k', pretrained=True)\n    if freeze_layers:\n        blocks = _get_feature_blocks(model)\n        _freeze_except_last_n(blocks, 2)\n    # handle different head naming\n    head = getattr(model, 'head', None)\n    print(head)\n    linear = getattr(head, 'linear')\n    in_features = 192\n    model.head.linear = nn.Sequential(\n        nn.Linear(in_features, 512),\n        nn.ReLU(inplace=True),\n        nn.Dropout(0.4),\n        nn.Linear(512, num_classes)\n    )\n    return model.to(device)\n    \ndef get_model_gernet(num_classes, freeze_layers=True, device='cuda'):\n    \"\"\"\n    Load and configure a GENet (General and Efficient Network) model with customizable classifier.\n    \n    Args:\n        num_classes: Number of output classes\n        freeze_layers: If True, freeze all but the last 2 blocks\n        device: Device to load the model on ('cuda' or 'cpu')\n        \n    Returns:\n        Configured GENet model\n    \"\"\"\n    model = timm.create_model('gernet_s.idstcv_in1k', pretrained=True)\n    \n    if freeze_layers:\n        # For GENet, we need to specifically handle its structure\n        # It typically has a 'stem' and 'stages' structure\n        if hasattr(model, 'stem') and hasattr(model, 'stages'):\n            # Freeze stem completely\n            for param in model.stem.parameters():\n                param.requires_grad = False\n                \n            # Freeze all stages except the last two\n            stages = list(model.stages.children())\n            total_stages = len(stages)\n            for i, stage in enumerate(stages):\n                requires_grad = (i >= total_stages - 2)\n                for param in stage.parameters():\n                    param.requires_grad = requires_grad\n        else:\n            # Fallback to generic approach\n            blocks = _get_feature_blocks(model)\n            _freeze_except_last_n(blocks, 2)\n    \n    # Replace classifier\n    in_features = model.head.fc.in_features\n    model.head.fc = nn.Sequential(\n        nn.Linear(in_features, 512),\n        nn.ReLU(inplace=True),\n        nn.Dropout(0.4),\n        nn.Linear(512, num_classes)\n    )\n    return model.to(device)\n\ndef get_model_regnetx(num_classes, freeze_layers=True, device='cuda'):\n    \"\"\"\n    Load and configure a RegNetX model with customizable classifier.\n    \n    Args:\n        num_classes: Number of output classes\n        freeze_layers: If True, freeze all but the last 2 blocks\n        device: Device to load the model on ('cuda' or 'cpu')\n        \n    Returns:\n        Configured RegNetX model\n    \"\"\"\n    model = timm.create_model('regnetx_008.tv2_in1k', pretrained=True)\n    \n    if freeze_layers:\n        # Looking at the error, we need to inspect the model structure carefully\n        # Print the model structure to understand it better in real use\n        # print(model)\n        \n        # Direct approach: check the model structure and freeze components individually\n        # First, freeze all parameters\n        for param in model.parameters():\n            param.requires_grad = False\n            \n        # Then unfreeze the last few layers manually based on RegNetX structure\n        # RegNetX typically has 'stem' + 'trunk' structure in timm\n        if hasattr(model, 'trunk'):\n            # Unfreeze final stages of the trunk\n            trunk_blocks = list(model.trunk.children())\n            # Unfreeze approximately last 25% of trunk blocks\n            unfreeze_from = max(0, int(len(trunk_blocks) * 0.75))\n            for i in range(unfreeze_from, len(trunk_blocks)):\n                for param in trunk_blocks[i].parameters():\n                    param.requires_grad = True\n        \n        # Always unfreeze the classifier/head for fine-tuning\n        for param in model.head.parameters():\n            param.requires_grad = True\n    \n    # Replace classifier\n    in_features = model.head.fc.in_features\n    model.head.fc = nn.Sequential(\n        nn.Linear(in_features, 512),\n        nn.ReLU(inplace=True),\n        nn.Dropout(0.4),\n        nn.Linear(512, num_classes)\n    )\n    return model.to(device)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-07T15:22:16.465842Z","iopub.execute_input":"2025-05-07T15:22:16.466084Z","iopub.status.idle":"2025-05-07T15:22:16.492719Z","shell.execute_reply.started":"2025-05-07T15:22:16.466067Z","shell.execute_reply":"2025-05-07T15:22:16.491782Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Training function","metadata":{}},{"cell_type":"code","source":"def train_model(model, criterion, optimizer, scheduler, train_loader, val_loader, early_stopping, epochs=15, use_ddp=False):\n    \"\"\"\n    Train the model and perform validation using multiple GPUs.\n    Supports both DataParallel (DP) and DistributedDataParallel (DDP) modes.\n    \n    Args:\n        model: Model to train\n        criterion: Loss function\n        optimizer: Optimizer for training\n        scheduler: Learning rate scheduler\n        train_loader: DataLoader for training data\n        val_loader: DataLoader for validation data\n        early_stopping: Early stopping handler\n        epochs: Maximum number of epochs to train\n        use_ddp: Whether to use DistributedDataParallel (True) or DataParallel (False)\n    \"\"\"\n    # Check available GPUs\n    num_gpus = torch.cuda.device_count()\n    if num_gpus < 2:\n        print(f\"Warning: Requested multi-GPU training but only {num_gpus} GPU(s) available. Continuing with available resources.\")\n    else:\n        print(f\"Using {num_gpus} GPUs for training\")\n    \n    # Setup device and model\n    if num_gpus >= 2:\n        if use_ddp:\n            # For DistributedDataParallel\n            import torch.distributed as dist\n            from torch.nn.parallel import DistributedDataParallel as DDP\n            \n            # Initialize process group\n            dist.init_process_group(backend='nccl')\n            local_rank = dist.get_rank()\n            torch.cuda.set_device(local_rank)\n            device = torch.device(f\"cuda:{local_rank}\")\n            \n            model = model.to(device)\n            model = DDP(model, device_ids=[local_rank])\n        else:\n            # For DataParallel (simpler to use)\n            device = torch.device(\"cuda:0\")\n            model = model.to(device)\n            model = torch.nn.DataParallel(model)\n    else:\n        # Single GPU\n        device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n        model = model.to(device)\n    \n    train_losses = []\n    val_losses = []\n    train_accs = []\n    val_accs = []\n    \n    # Store validation predictions and labels for final evaluation\n    all_val_labels = []\n    all_val_preds = []\n    all_val_scores = []\n    \n    for epoch in range(epochs):\n        print(f\"Epoch {epoch+1}/{epochs}\")\n        \n        # Training phase\n        model.train()\n        running_loss = 0.0\n        correct = 0\n        total = 0\n        \n        for inputs, labels in tqdm(train_loader, desc=\"Training\"):\n            inputs, labels = inputs.to(device), labels.to(device)\n            \n            optimizer.zero_grad()\n            outputs = model(inputs)\n            loss = criterion(outputs, labels)\n            loss.backward()\n            optimizer.step()\n            \n            running_loss += loss.item() * inputs.size(0)\n            _, predicted = torch.max(outputs, 1)\n            total += labels.size(0)\n            correct += (predicted == labels).sum().item()\n        \n        epoch_train_loss = running_loss / len(train_loader.dataset)\n        epoch_train_acc = correct / total\n        train_losses.append(epoch_train_loss)\n        train_accs.append(epoch_train_acc)\n        \n        # Validation phase\n        model.eval()\n        running_loss = 0.0\n        correct = 0\n        total = 0\n        \n        all_labels = []\n        all_preds = []\n        all_scores = []\n        \n        with torch.no_grad():\n            for inputs, labels in tqdm(val_loader, desc=\"Validation\"):\n                inputs, labels = inputs.to(device), labels.to(device)\n                outputs = model(inputs)\n                loss = criterion(outputs, labels)\n                \n                running_loss += loss.item() * inputs.size(0)\n                probs = F.softmax(outputs, dim=1)\n                _, predicted = torch.max(outputs, 1)\n                total += labels.size(0)\n                correct += (predicted == labels).sum().item()\n                \n                all_labels.extend(labels.cpu().numpy().tolist())\n                all_preds.extend(predicted.cpu().numpy().tolist())\n                all_scores.append(probs.cpu().numpy())\n        \n        epoch_val_loss = running_loss / len(val_loader.dataset)\n        epoch_val_acc = correct / total\n        val_losses.append(epoch_val_loss)\n        val_accs.append(epoch_val_acc)\n        \n        all_scores = np.vstack(all_scores) if all_scores else np.array([])\n        \n        # Store validation results for the final epoch\n        all_val_labels = all_labels\n        all_val_preds = all_preds\n        all_val_scores = all_scores\n        \n        # Update learning rate scheduler\n        scheduler.step(epoch_val_loss)\n        \n        print(f\"Train Loss: {epoch_train_loss:.4f}, Train Acc: {epoch_train_acc:.4f}\")\n        print(f\"Val Loss: {epoch_val_loss:.4f}, Val Acc: {epoch_val_acc:.4f}\")\n        print(f\"Learning rate: {optimizer.param_groups[0]['lr']:.6f}\")\n        \n        # Check early stopping\n        early_stopping(epoch_val_loss)\n        if early_stopping.early_stop:\n            print(\"Early stopping triggered!\")\n            break\n            \n        # Free up memory\n        del all_labels, all_preds, all_scores\n        gc.collect()\n        torch.cuda.empty_cache()\n    \n    # Clean up DDP if used\n    if num_gpus >= 2 and use_ddp:\n        dist.destroy_process_group()\n    \n    return model, train_losses, val_losses, train_accs, val_accs, all_val_labels, all_val_preds, all_val_scores\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-07T15:22:16.494736Z","iopub.execute_input":"2025-05-07T15:22:16.495264Z","iopub.status.idle":"2025-05-07T15:22:16.517964Z","shell.execute_reply.started":"2025-05-07T15:22:16.495245Z","shell.execute_reply":"2025-05-07T15:22:16.517204Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Evaluation plotting functions","metadata":{}},{"cell_type":"code","source":"def plot_roc_curves(y_true, y_scores, class_names):\n    \"\"\"\n    Plot ROC curves for multi-class classification.\n    \n    Parameters:\n    - y_true: true labels\n    - y_scores: predicted probability scores from model\n    - class_names: list of class names\n    \"\"\"\n    # Ensure inputs are numpy arrays\n    if torch.is_tensor(y_true):\n        y_true = y_true.cpu().numpy()\n    if torch.is_tensor(y_scores):\n        y_scores = y_scores.cpu().numpy()\n    \n    n_classes = len(class_names)\n    \n    # Binarize the labels for one-vs-rest ROC calculation\n    y_true_bin = label_binarize(y_true, classes=range(n_classes))\n    \n    # Compute ROC curve and ROC area for each class\n    fpr = {}\n    tpr = {}\n    roc_auc = {}\n    \n    plt.figure(figsize=(12, 8))\n    \n    for i in range(n_classes):\n        fpr[i], tpr[i], _ = roc_curve(y_true_bin[:, i], y_scores[:, i])\n        roc_auc[i] = auc(fpr[i], tpr[i])\n        \n        plt.plot(fpr[i], tpr[i], lw=2,\n                 label=f'{class_names[i]} (area = {roc_auc[i]:.2f})')\n    \n    # Plot the diagonal (random classifier)\n    plt.plot([0, 1], [0, 1], 'k--', lw=2)\n    \n    # Calculate and plot micro-average ROC curve\n    fpr[\"micro\"], tpr[\"micro\"], _ = roc_curve(y_true_bin.ravel(), y_scores.ravel())\n    roc_auc[\"micro\"] = auc(fpr[\"micro\"], tpr[\"micro\"])\n    plt.plot(fpr[\"micro\"], tpr[\"micro\"], \n             label=f'Micro-average (area = {roc_auc[\"micro\"]:.2f})', \n             lw=2, linestyle=':', color='deeppink')\n    \n    plt.xlim([0.0, 1.0])\n    plt.ylim([0.0, 1.05])\n    plt.xlabel('False Positive Rate')\n    plt.ylabel('True Positive Rate')\n    plt.title('ROC Curves')\n    plt.legend(loc=\"lower right\")\n    plt.grid(True, alpha=0.3)\n    plt.tight_layout()\n    plt.show()\n    \n    # Return the AUC values for reporting\n    return roc_auc\n\ndef plot_pr_curves(y_true, y_scores, class_names):\n    \"\"\"\n    Plot Precision-Recall curves for multi-class classification.\n    \n    Parameters:\n    - y_true: true labels\n    - y_scores: predicted probability scores from model\n    - class_names: list of class names\n    \"\"\"\n    # Ensure inputs are numpy arrays\n    if torch.is_tensor(y_true):\n        y_true = y_true.cpu().numpy()\n    if torch.is_tensor(y_scores):\n        y_scores = y_scores.cpu().numpy()\n    \n    n_classes = len(class_names)\n    \n    # Binarize the labels\n    y_true_bin = label_binarize(y_true, classes=range(n_classes))\n    \n    # Compute PR curve and average precision for each class\n    precision = {}\n    recall = {}\n    avg_precision = {}\n    \n    plt.figure(figsize=(12, 8))\n    \n    for i in range(n_classes):\n        precision[i], recall[i], _ = precision_recall_curve(y_true_bin[:, i], y_scores[:, i])\n        avg_precision[i] = average_precision_score(y_true_bin[:, i], y_scores[:, i])\n        \n        plt.plot(recall[i], precision[i], lw=2,\n                 label=f'{class_names[i]} (AP = {avg_precision[i]:.2f})')\n    \n    # Calculate and plot micro-average PR curve\n    precision[\"micro\"], recall[\"micro\"], _ = precision_recall_curve(\n        y_true_bin.ravel(), y_scores.ravel())\n    avg_precision[\"micro\"] = average_precision_score(y_true_bin.ravel(), y_scores.ravel())\n    \n    plt.plot(recall[\"micro\"], precision[\"micro\"],\n             label=f'Micro-average (AP = {avg_precision[\"micro\"]:.2f})',\n             lw=2, linestyle=':', color='deeppink')\n    \n    plt.xlim([0.0, 1.0])\n    plt.ylim([0.0, 1.05])\n    plt.xlabel('Recall')\n    plt.ylabel('Precision')\n    plt.title('Precision-Recall Curves')\n    plt.legend(loc=\"best\")\n    plt.grid(True, alpha=0.3)\n    plt.tight_layout()\n    plt.show()\n    \n    # Return the average precision values for reporting\n    return avg_precision\n\ndef plot_accuracy_and_loss(train_losses, val_losses, train_accs, val_accs):\n    plt.figure(figsize=(12, 5))\n    # Accuracy curve\n    plt.subplot(1, 2, 1)\n    plt.plot(train_accs, label=\"Train Accuracy\")\n    plt.plot(val_accs, label=\"Validation Accuracy\")\n    plt.xlabel(\"Epochs\")\n    plt.ylabel(\"Accuracy\")\n    plt.title(\"Accuracy Curve\")\n    plt.legend()\n    plt.grid(True)\n    \n    # Loss curve\n    plt.subplot(1, 2, 2)\n    plt.plot(train_losses, label=\"Train Loss\")\n    plt.plot(val_losses, label=\"Validation Loss\")\n    plt.xlabel(\"Epochs\")\n    plt.ylabel(\"Loss\")\n    plt.title(\"Loss Curve\")\n    plt.legend()\n    plt.grid(True)\n    \n    plt.tight_layout()\n    plt.show()\n\ndef plot_confusion_matrix(y_true, y_pred, class_names):\n    # Ensure we're working with numpy arrays\n    y_true = np.array(y_true)\n    y_pred = np.array(y_pred)\n    \n    # Get unique values in both arrays\n    unique_values = np.unique(np.concatenate([y_true, y_pred]))\n    print(f\"Unique values in confusion matrix data: {unique_values}\")\n    \n    # Create the confusion matrix with explicit labels\n    cm = confusion_matrix(y_true, y_pred, labels=range(len(class_names)))\n    \n    plt.figure(figsize=(10, 8))\n    sns.heatmap(cm, annot=True, fmt=\"d\", cmap=\"Blues\",\n                xticklabels=class_names, yticklabels=class_names)\n    plt.title(\"Confusion Matrix\")\n    plt.xlabel(\"Predicted\")\n    plt.ylabel(\"True\")\n    plt.tight_layout()\n    plt.show()\n\ndef plot_per_class_accuracy(y_true, y_pred, class_names):\n    # Convert to numpy arrays\n    y_true = np.array(y_true)\n    y_pred = np.array(y_pred)\n    \n    # Get number of expected classes\n    num_classes = len(class_names)\n    \n    # Create the confusion matrix with explicit labels\n    cm = confusion_matrix(y_true, y_pred, labels=range(num_classes))\n    \n    # Calculate per-class accuracy\n    per_class_accuracy = np.zeros(num_classes)\n    for i in range(num_classes):\n        if i < cm.shape[0] and np.sum(cm[i, :]) > 0:\n            per_class_accuracy[i] = cm[i, i] / np.sum(cm[i, :])\n    \n    # Create the bar plot\n    plt.figure(figsize=(14, 7))\n    plt.bar(range(num_classes), per_class_accuracy, color=\"skyblue\")\n    plt.xticks(range(num_classes), class_names, rotation=45, ha='right')\n    plt.xlabel(\"Classes\")\n    plt.ylabel(\"Accuracy\")\n    plt.title(\"Per-Class Accuracy\")\n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-07T15:22:16.518997Z","iopub.execute_input":"2025-05-07T15:22:16.519284Z","iopub.status.idle":"2025-05-07T15:22:16.546032Z","shell.execute_reply.started":"2025-05-07T15:22:16.519262Z","shell.execute_reply":"2025-05-07T15:22:16.545330Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from sklearn.metrics import cohen_kappa_score\n\ndef compute_classification_metrics(y_true, y_pred, y_scores, num_classes, class_names, model_name=\"\"):\n    \"\"\"\n    Compute comprehensive classification metrics including ROC AUC, PR AUC, and Cohen's Kappa.\n    \n    Parameters:\n    - y_true: true labels\n    - y_pred: predicted labels\n    - y_scores: predicted probability scores from model\n    - num_classes: number of classes\n    - class_names: list of class names\n    - model_name: name of the model (for display purposes)\n    \n    Returns:\n    - accuracy: overall accuracy score\n    - report_dict: classification report as dictionary\n    - roc_auc_dict: ROC AUC scores by class\n    - pr_auc_dict: PR AUC scores by class\n    - kappa: Cohen's Kappa score\n    \"\"\"\n    # Calculate accuracy\n    accuracy = accuracy_score(y_true, y_pred)\n    print(f\"Overall Accuracy: {accuracy:.4f}\")\n    \n    # Calculate and display Cohen's Kappa\n    kappa = cohen_kappa_score(y_true, y_pred)\n    print(f\"Cohen's Kappa Score: {kappa:.4f}\")\n    \n    # Generate classification report\n    report = classification_report(y_true, y_pred, target_names=class_names, output_dict=True)\n    \n    # Print formatted classification report\n    print(\"\\nClassification Report:\")\n    print(classification_report(y_true, y_pred, target_names=class_names))\n    \n    # Calculate ROC curves and AUC for each class\n    print(\"\\nCalculating ROC curves...\")\n    roc_auc_dict = plot_roc_curves(y_true, y_scores, class_names)\n    \n    # Calculate PR curves and AUC for each class\n    print(\"\\nCalculating PR curves...\")\n    pr_auc_dict = plot_pr_curves(y_true, y_scores, class_names)\n    \n    # Return metrics for comparison\n    return accuracy, report, roc_auc_dict, pr_auc_dict, kappa\n\n# Also update evaluate_on_test_set to include kappa\ndef evaluate_on_test_set(model, test_loader, dataset):\n    \"\"\"Evaluate a trained model on test dataset\"\"\"\n    class_names = dataset.classes\n    num_classes = len(class_names)\n    \n    model.eval()\n    device = next(model.parameters()).device\n    \n    all_labels = []\n    all_preds = []\n    all_scores = []\n    \n    with torch.no_grad():\n        for inputs, labels in test_loader:\n            inputs, labels = inputs.to(device), labels.to(device)\n            outputs = model(inputs)\n            _, preds = torch.max(outputs, 1)\n            \n            all_labels.extend(labels.cpu().numpy())\n            all_preds.extend(preds.cpu().numpy())\n            all_scores.append(torch.nn.functional.softmax(outputs, dim=1).cpu().numpy())\n    \n    all_scores = np.vstack(all_scores)\n    \n    # Compute metrics including kappa\n    accuracy, report_dict, roc_auc_dict, pr_auc_dict, kappa = compute_classification_metrics(\n        all_labels, all_preds, all_scores, num_classes, class_names)\n    \n    # Build results dictionary with kappa\n    results = {\n        'accuracy': accuracy,\n        'report': report_dict,\n        'roc_auc': roc_auc_dict,\n        'pr_auc': pr_auc_dict,\n        'kappa': kappa\n    }\n    \n    return results","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-07T15:22:16.547383Z","iopub.execute_input":"2025-05-07T15:22:16.547662Z","iopub.status.idle":"2025-05-07T15:22:16.571142Z","shell.execute_reply.started":"2025-05-07T15:22:16.547638Z","shell.execute_reply":"2025-05-07T15:22:16.570633Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Main training function","metadata":{}},{"cell_type":"code","source":"# Update the model_train function to include kappa in the results\ndef model_train(model, train_loader, val_loader, dataset, epochs=20):\n    model_name = type(model).__name__\n    if hasattr(model, 'pretrained_cfg') and 'name' in model.pretrained_cfg:\n        model_name = model.pretrained_cfg['name']\n        \n    print(f\"\\n{'='*20} Training {model_name} {'='*20}\\n\")\n        \n    class_names = dataset.classes\n    num_classes = len(class_names)\n    learning_rate = 0.001\n    \n    try:\n        optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=3)\n        early_stopping = EarlyStopping(patience=5)\n        \n        model, train_losses, val_losses, train_accs, val_accs, val_labels, val_preds, val_scores = train_model(\n            model, nn.CrossEntropyLoss(), optimizer, scheduler,\n            train_loader, val_loader, early_stopping, epochs=epochs, use_ddp=False\n        )\n        \n        print(f\"\\n{'='*20} Evaluation for {model_name} {'='*20}\\n\")\n        \n        # Plot training curves\n        plot_accuracy_and_loss(train_losses, val_losses, train_accs, val_accs)\n        \n        # Process validation predictions and labels\n        try:\n            plot_confusion_matrix(val_labels, val_preds, class_names)\n            plot_per_class_accuracy(val_labels, val_preds, class_names)\n            \n            # Get metrics from the updated function including kappa\n            accuracy, report_dict, roc_auc_dict, pr_auc_dict, kappa = compute_classification_metrics(\n                val_labels, val_preds, val_scores, num_classes, class_names, model_name)\n            \n            # Build a results dictionary including kappa\n            results = {\n                'accuracy': accuracy,\n                'report': report_dict,\n                'roc_auc': roc_auc_dict,\n                'pr_auc': pr_auc_dict,\n                'kappa': kappa\n            }\n            \n            return results\n        except Exception as viz_error:\n            print(f\"Error in visualization: {viz_error}\")\n            import traceback\n            traceback.print_exc()\n            return {'accuracy': None}\n            \n    except Exception as e:\n        print(f'Error occurred when training {model_name}: {e}')\n        import traceback\n        traceback.print_exc()\n        return {'accuracy': None}\n    finally:\n        # Clean up memory\n        if 'optimizer' in locals():\n            del optimizer\n        if 'scheduler' in locals():\n            del scheduler\n        if 'early_stopping' in locals():\n            del early_stopping\n        if 'train_losses' in locals():\n            del train_losses\n            del val_losses\n            del train_accs\n            del val_accs\n            del val_labels\n            del val_preds\n            del val_scores\n        \n        gc.collect()\n        torch.cuda.empty_cache()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-07T15:22:16.571806Z","iopub.execute_input":"2025-05-07T15:22:16.572176Z","iopub.status.idle":"2025-05-07T15:22:16.594545Z","shell.execute_reply.started":"2025-05-07T15:22:16.572146Z","shell.execute_reply":"2025-05-07T15:22:16.593850Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Comparison Function","metadata":{}},{"cell_type":"code","source":"def compare_models(models, train_loader, val_loader, test_loader, dataset, epochs=20, names=None):\n    if names is None:\n        names = [f\"Model {i+1}\" for i in range(len(models))]\n    \n    val_results = {}\n    test_results = {}\n    best_model_obj = None\n    best_accuracy = -1\n    best_model_name = \"\"\n    \n    # Summary dictionaries for metrics\n    val_roc_auc_summary = {}\n    test_roc_auc_summary = {}\n    val_pr_auc_summary = {}\n    test_pr_auc_summary = {}\n    val_kappa_summary = {}\n    test_kappa_summary = {}\n    \n    for i, (model, name) in enumerate(zip(models, names)):\n        print(f\"\\n\\n{'#'*30} Training {name} ({i+1}/{len(models)}) {'#'*30}\\n\")\n        model_results = model_train(model, train_loader, val_loader, dataset, epochs)\n        \n        # Extract accuracy from results\n        accuracy = model_results.get('accuracy')\n        val_results[name] = accuracy\n        \n        # Extract and store metrics\n        if 'roc_auc' in model_results and 'micro' in model_results['roc_auc']:\n            val_roc_auc_summary[name] = model_results['roc_auc']['micro']\n        else:\n            val_roc_auc_summary[name] = None\n            \n        if 'pr_auc' in model_results and 'micro' in model_results['pr_auc']:\n            val_pr_auc_summary[name] = model_results['pr_auc']['micro']\n        else:\n            val_pr_auc_summary[name] = None\n            \n        # Store kappa score\n        if 'kappa' in model_results:\n            val_kappa_summary[name] = model_results['kappa']\n        else:\n            val_kappa_summary[name] = None\n        \n        # Evaluate on test set\n        if accuracy is not None:\n            print(f\"\\n{'='*20} Testing {name} on Test Set {'='*20}\\n\")\n            test_model_results = evaluate_on_test_set(model, test_loader, dataset)\n            \n            # Extract accuracy from test results\n            test_accuracy = test_model_results.get('accuracy')\n            test_results[name] = test_accuracy\n            \n            # Extract and store test metrics\n            if 'roc_auc' in test_model_results and 'micro' in test_model_results['roc_auc']:\n                test_roc_auc_summary[name] = test_model_results['roc_auc']['micro']\n            else:\n                test_roc_auc_summary[name] = None\n                \n            if 'pr_auc' in test_model_results and 'micro' in test_model_results['pr_auc']:\n                test_pr_auc_summary[name] = test_model_results['pr_auc']['micro']\n            else:\n                test_pr_auc_summary[name] = None\n                \n            # Store test kappa score\n            if 'kappa' in test_model_results:\n                test_kappa_summary[name] = test_model_results['kappa']\n            else:\n                test_kappa_summary[name] = None\n            \n            # Track best model\n            if test_accuracy > best_accuracy:\n                best_accuracy = test_accuracy\n                best_model_obj = copy.deepcopy(model)\n                best_model_name = name\n    \n    # Print comprehensive comparison\n    print(\"\\n\\n\" + \"=\"*100)\n    print(\"COMPREHENSIVE MODEL COMPARISON\")\n    print(\"=\"*100)\n    print(f\"{'Model':<20}{'Val Acc':<10}{'Test Acc':<10}{'Val ROC AUC':<14}{'Test ROC AUC':<14}{'Val PR AUC':<14}{'Test PR AUC':<14}{'Val Kappa':<14}{'Test Kappa':<14}\")\n    print(\"-\"*100)\n    \n    for name in val_results.keys():\n        val_acc = val_results[name]\n        test_acc = test_results.get(name, None)\n        val_roc = val_roc_auc_summary.get(name, None)\n        test_roc = test_roc_auc_summary.get(name, None)\n        val_pr = val_pr_auc_summary.get(name, None)\n        test_pr = test_pr_auc_summary.get(name, None)\n        val_kappa = val_kappa_summary.get(name, None)\n        test_kappa = test_kappa_summary.get(name, None)\n        \n        # Format values for display\n        val_acc_str = f\"{val_acc:.4f}\" if val_acc is not None else \"Failed\"\n        test_acc_str = f\"{test_acc:.4f}\" if test_acc is not None else \"N/A\"\n        val_roc_str = f\"{val_roc:.4f}\" if val_roc is not None else \"N/A\"\n        test_roc_str = f\"{test_roc:.4f}\" if test_roc is not None else \"N/A\"\n        val_pr_str = f\"{val_pr:.4f}\" if val_pr is not None else \"N/A\"\n        test_pr_str = f\"{test_pr:.4f}\" if test_pr is not None else \"N/A\"\n        val_kappa_str = f\"{val_kappa:.4f}\" if val_kappa is not None else \"N/A\"\n        test_kappa_str = f\"{test_kappa:.4f}\" if test_kappa is not None else \"N/A\"\n            \n        print(f\"{name:<20}{val_acc_str:<10}{test_acc_str:<10}{val_roc_str:<14}{test_roc_str:<14}{val_pr_str:<14}{test_pr_str:<14}{val_kappa_str:<14}{test_kappa_str:<14}\")\n    \n    # Identify best model based on test metrics\n    if test_results:\n        # Best model by accuracy\n        best_acc_model = max(test_results.items(), key=lambda x: x[1] if x[1] is not None else -1)\n        print(f\"\\nBest model by accuracy: {best_acc_model[0]} (Test Accuracy: {best_acc_model[1]:.4f})\")\n        \n        # Best model by ROC AUC (if available)\n        if any(v is not None for v in test_roc_auc_summary.values()):\n            best_roc_model = max(\n                [(k, v) for k, v in test_roc_auc_summary.items() if v is not None], \n                key=lambda x: x[1] if x[1] is not None else -1\n            )\n            print(f\"Best model by ROC AUC: {best_roc_model[0]} (Test ROC AUC: {best_roc_model[1]:.4f})\")\n        \n        # Best model by PR AUC (if available)\n        if any(v is not None for v in test_pr_auc_summary.values()):\n            best_pr_model = max(\n                [(k, v) for k, v in test_pr_auc_summary.items() if v is not None], \n                key=lambda x: x[1] if x[1] is not None else -1\n            )\n            print(f\"Best model by PR AUC: {best_pr_model[0]} (Test PR AUC: {best_pr_model[1]:.4f})\")\n            \n        # Best model by Kappa (if available)\n        if any(v is not None for v in test_kappa_summary.values()):\n            best_kappa_model = max(\n                [(k, v) for k, v in test_kappa_summary.items() if v is not None], \n                key=lambda x: x[1] if x[1] is not None else -1\n            )\n            print(f\"Best model by Cohen's Kappa: {best_kappa_model[0]} (Test Kappa: {best_kappa_model[1]:.4f})\")\n        \n        # Save the best model (by accuracy)\n        if best_model_obj is not None:\n            try:\n                model_save_path = f\"best_model_{best_model_name.lower().replace(' ', '_')}.pth\"\n                torch.save(best_model_obj.state_dict(), model_save_path)\n                print(f\"Best model saved to {model_save_path}\")\n            except Exception as save_error:\n                print(f\"Error saving best model: {save_error}\")\n    else:\n        print(\"\\nNo models successfully completed testing.\")\n    \n    print(\"=\"*100)\n    \n    # Visualize comparison\n    try:\n        # Create bar charts comparing different metrics\n        plot_model_comparison(val_results, test_results, val_roc_auc_summary, \n                             test_roc_auc_summary, val_pr_auc_summary, test_pr_auc_summary,\n                             val_kappa_summary, test_kappa_summary)\n    except Exception as viz_error:\n        print(f\"Error in comparison visualization: {viz_error}\")\n        import traceback\n        traceback.print_exc()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-07T15:22:16.595427Z","iopub.execute_input":"2025-05-07T15:22:16.595661Z","iopub.status.idle":"2025-05-07T15:22:16.620099Z","shell.execute_reply.started":"2025-05-07T15:22:16.595635Z","shell.execute_reply":"2025-05-07T15:22:16.619560Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def plot_model_comparison(val_acc, test_acc, val_roc, test_roc, val_pr, test_pr, val_kappa, test_kappa):\n    \"\"\"\n    Create visualizations to compare model performance across different metrics including Cohen's Kappa.\n    \"\"\"\n    # Get the list of model names (should be the same across all dictionaries)\n    models = list(val_acc.keys())\n    \n    # Create a figure with 4 subplots for Accuracy, ROC AUC, PR AUC, and Kappa\n    fig, (ax1, ax2, ax3, ax4) = plt.subplots(4, 1, figsize=(12, 24))\n    \n    # Plot Accuracy\n    x = np.arange(len(models))\n    width = 0.35\n    \n    val_acc_values = [val_acc.get(model, None) for model in models]\n    test_acc_values = [test_acc.get(model, None) for model in models]\n    \n    # Replace None with NaN for plotting\n    val_acc_values = [v if v is not None else float('nan') for v in val_acc_values]\n    test_acc_values = [v if v is not None else float('nan') for v in test_acc_values]\n    \n    ax1.bar(x - width/2, val_acc_values, width, label='Validation', color='skyblue')\n    ax1.bar(x + width/2, test_acc_values, width, label='Test', color='salmon')\n    ax1.set_ylabel('Accuracy')\n    ax1.set_title('Model Accuracy Comparison')\n    ax1.set_xticks(x)\n    ax1.set_xticklabels(models, rotation=45, ha='right')\n    ax1.legend()\n    ax1.grid(True, alpha=0.3)\n    \n    # Plot ROC AUC\n    val_roc_values = [val_roc.get(model, None) for model in models]\n    test_roc_values = [test_roc.get(model, None) for model in models]\n    \n    # Replace None with NaN for plotting\n    val_roc_values = [v if v is not None else float('nan') for v in val_roc_values]\n    test_roc_values = [v if v is not None else float('nan') for v in test_roc_values]\n    \n    ax2.bar(x - width/2, val_roc_values, width, label='Validation', color='skyblue')\n    ax2.bar(x + width/2, test_roc_values, width, label='Test', color='salmon')\n    ax2.set_ylabel('ROC AUC')\n    ax2.set_title('Model ROC AUC Comparison')\n    ax2.set_xticks(x)\n    ax2.set_xticklabels(models, rotation=45, ha='right')\n    ax2.legend()\n    ax2.grid(True, alpha=0.3)\n    \n    # Plot PR AUC\n    val_pr_values = [val_pr.get(model, None) for model in models]\n    test_pr_values = [test_pr.get(model, None) for model in models]\n    \n    # Replace None with NaN for plotting\n    val_pr_values = [v if v is not None else float('nan') for v in val_pr_values]\n    test_pr_values = [v if v is not None else float('nan') for v in test_pr_values]\n    \n    ax3.bar(x - width/2, val_pr_values, width, label='Validation', color='skyblue')\n    ax3.bar(x + width/2, test_pr_values, width, label='Test', color='salmon')\n    ax3.set_ylabel('PR AUC')\n    ax3.set_title('Model PR AUC Comparison')\n    ax3.set_xticks(x)\n    ax3.set_xticklabels(models, rotation=45, ha='right')\n    ax3.legend()\n    ax3.grid(True, alpha=0.3)\n    \n    # Plot Kappa scores\n    val_kappa_values = [val_kappa.get(model, None) for model in models]\n    test_kappa_values = [test_kappa.get(model, None) for model in models]\n    \n    # Replace None with NaN for plotting\n    val_kappa_values = [v if v is not None else float('nan') for v in val_kappa_values]\n    test_kappa_values = [v if v is not None else float('nan') for v in test_kappa_values]\n    \n    ax4.bar(x - width/2, val_kappa_values, width, label='Validation', color='skyblue')\n    ax4.bar(x + width/2, test_kappa_values, width, label='Test', color='salmon')\n    ax4.set_ylabel(\"Cohen's Kappa\")\n    ax4.set_title(\"Model Cohen's Kappa Comparison\")\n    ax4.set_xticks(x)\n    ax4.set_xticklabels(models, rotation=45, ha='right')\n    ax4.legend()\n    ax4.grid(True, alpha=0.3)\n    \n    plt.tight_layout()\n    plt.show()\n    \n    # Create a comprehensive heatmap for all metrics\n    try:\n        plot_metrics_heatmap(models, val_acc_values, test_acc_values, \n                            val_roc_values, test_roc_values,\n                            val_pr_values, test_pr_values,\n                            val_kappa_values, test_kappa_values)\n    except Exception as e:\n        print(f\"Error creating metrics heatmap: {e}\")\n\ndef plot_metrics_heatmap(models, val_acc, test_acc, val_roc, test_roc, val_pr, test_pr, val_kappa, test_kappa):\n    \"\"\"\n    Create a heatmap visualization of all metrics for easy comparison across models.\n    \"\"\"\n    # Prepare data for heatmap\n    metric_names = ['Val Acc', 'Test Acc', 'Val ROC', 'Test ROC', \n                   'Val PR', 'Test PR', 'Val Kappa', 'Test Kappa']\n    \n    data = np.array([\n        val_acc, test_acc, val_roc, test_roc, val_pr, test_pr, val_kappa, test_kappa\n    ])\n    \n    # Create the heatmap\n    plt.figure(figsize=(12, 8))\n    ax = sns.heatmap(data, annot=True, fmt=\".4f\", cmap=\"YlGnBu\", \n                    xticklabels=models, yticklabels=metric_names)\n    \n    plt.title(\"Comprehensive Model Performance Metrics\")\n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-07T15:22:16.595427Z","iopub.execute_input":"2025-05-07T15:22:16.595661Z","iopub.status.idle":"2025-05-07T15:22:16.620099Z","shell.execute_reply.started":"2025-05-07T15:22:16.595635Z","shell.execute_reply":"2025-05-07T15:22:16.619560Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Main Function","metadata":{}},{"cell_type":"code","source":"train_dir = '/kaggle/input/eye-disease-image-dataset/Augmented Dataset/Augmented Dataset'  # For training (pre-augmented data)\neval_dir = '/kaggle/input/eye-disease-image-dataset/Original Dataset/Original Dataset'  # For val and test\nepochs = 15\nclasses_to_exclude = [\"Pterygium\"]\nbatch_size = 32\n\nseed = 42\nrandom.seed(seed)\nnp.random.seed(seed)\ntorch.manual_seed(seed)\nif torch.cuda.is_available():\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n\n# Define transformations\ntransform = transforms.Compose([\n    transforms.Resize(256),\n    transforms.CenterCrop(224),\n    transforms.ToTensor(),\n    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n])\n\ntry:\n    # Load datasets\n    print(f\"Loading training dataset from {train_dir}...\")\n    train_dataset = datasets.ImageFolder(root=train_dir, transform=transform)\n    \n    print(f\"Loading evaluation dataset from {eval_dir}...\")\n    eval_dataset = datasets.ImageFolder(root=eval_dir, transform=transform)\n    \n    # Print dataset information\n    print(f\"Training dataset loaded with {len(train_dataset)} images across {len(train_dataset.classes)} classes.\")\n    print(f\"Training classes: {train_dataset.classes}\")\n    \n    print(f\"Evaluation dataset loaded with {len(eval_dataset)} images across {len(eval_dataset.classes)} classes.\")\n    print(f\"Evaluation classes: {eval_dataset.classes}\")\n    \n    # Filter datasets if needed\n    excluded_classes = classes_to_exclude or []\n    if excluded_classes:\n        print(f\"Filtering out classes: {excluded_classes}\")\n        filtered_train_dataset = FilteredImageDataset(train_dataset, excluded_classes=excluded_classes)\n        filtered_eval_dataset = FilteredImageDataset(eval_dataset, excluded_classes=excluded_classes)\n    else:\n        filtered_train_dataset = train_dataset\n        filtered_eval_dataset = eval_dataset\n    \n    # Check if the filtered classes match between training and evaluation datasets\n    if set(filtered_train_dataset.classes) != set(filtered_eval_dataset.classes):\n        print(\"Warning: Class mismatch between filtered training and evaluation datasets!\")\n        print(f\"Filtered training classes: {filtered_train_dataset.classes}\")\n        print(f\"Filtered evaluation classes: {filtered_eval_dataset.classes}\")\n        \n        # Find common classes\n        common_classes = set(filtered_train_dataset.classes).intersection(set(filtered_eval_dataset.classes))\n        print(f\"Common classes: {common_classes}\")\n        \n        # Create additional filtering based on common classes\n        filtered_train_dataset = FilteredImageDataset(train_dataset, \n                                                   included_classes=common_classes)\n        filtered_eval_dataset = FilteredImageDataset(eval_dataset, \n                                                  included_classes=common_classes)\n    \n    # Split evaluation dataset into validation and test sets\n    eval_ratio = 0.7  # 70% validation, 30% test\n    eval_size = len(filtered_eval_dataset)\n    val_size = int(eval_ratio * eval_size)\n    test_size = eval_size - val_size\n    \n    val_dataset, test_dataset = random_split(filtered_eval_dataset, [val_size, test_size])\n    \n    print(f\"Training set size: {len(filtered_train_dataset)}\")\n    print(f\"Validation set size: {len(val_dataset)}\")\n    print(f\"Test set size: {len(test_dataset)}\")\n    \n    # Create data loaders\n    train_loader = DataLoader(filtered_train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)\n    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4)\n    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)\n    \n    # Get the number of classes (after filtering)\n    num_classes = len(filtered_train_dataset.classes)\n    print(f\"Number of classes after filtering: {num_classes}\")\n    print(f\"Classes after filtering: {filtered_train_dataset.classes}\")\n    \n    # Initialize models\n    print(\"Initializing models...\")\n    \n    # Initialize models with the updated number of classes\n    all_models = {\n        \"MobileNetV4\": get_model_mobilenetv4(num_classes, freeze_layers=True),\n        \"LeViT\": get_model_levit(num_classes, freeze_layers=True),\n        \"EfficientViT\": get_model_efficientvit(num_classes, freeze_layers=True),\n        \"GENet\": get_model_gernet(num_classes, freeze_layers=True),\n        \"RegNetX\": get_model_regnetx(num_classes, freeze_layers=True)\n    }\n    \n    models = list(all_models.values())\n    model_names = list(all_models.keys())\n    \n    # Train and compare models\n    print(\"Starting model training and comparison...\")\n    compare_models(models, train_loader, val_loader, test_loader, filtered_train_dataset, epochs=epochs, names=model_names)\n    \nexcept Exception as e:\n    print(f\"Error in eye disease classification pipeline: {e}\")\n    import traceback\n    traceback.print_exc()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-07T15:22:59.764465Z","iopub.execute_input":"2025-05-07T15:22:59.765257Z","execution_failed":"2025-05-07T15:22:59.047Z"}},"outputs":[],"execution_count":null}]}

utils/Callback.py

@@ -0,0 +1,21 @@
+# Early stopping class
+class EarlyStopping:
+    def __init__(self, patience=5, delta=0):
+        self.patience = patience
+        self.delta = delta
+        self.counter = 0
+        self.best_score = None
+        self.early_stop = False
+
+    def __call__(self, val_loss):
+        score = -val_loss
+
+        if self.best_score is None:
+            self.best_score = score
+        elif score < self.best_score + self.delta:
+            self.counter += 1
+            if self.counter >= self.patience:
+                self.early_stop = True
+        else:
+            self.best_score = score
+            self.counter = 0

utils/Comparator.py

@@ -0,0 +1,179 @@
+from Trainer import model_train
+from Evaluator import ClassificationEvaluator
+
+
+def compare_models(
+    models, train_loader, val_loader, test_loader, dataset, epochs=20, names=None
+):
+    if names is None:
+        names = [f"Model {i+1}" for i in range(len(models))]
+
+    val_results = {}
+    test_results = {}
+    best_model_obj = None
+    best_accuracy = -1
+    best_model_name = ""
+
+    # Summary dictionaries for metrics
+    val_roc_auc_summary = {}
+    test_roc_auc_summary = {}
+    val_pr_auc_summary = {}
+    test_pr_auc_summary = {}
+    val_kappa_summary = {}
+    test_kappa_summary = {}
+
+    for i, (model, name) in enumerate(zip(models, names)):
+        evaluator = ClassificationEvaluator(
+            num_classes=len(dataset.classes),
+            class_names=dataset.classes,
+        )
+
+        print(f"\n\n{'#'*30} Training {name} ({i+1}/{len(models)}) {'#'*30}\n")
+        model_results = model_train(model, train_loader, val_loader, dataset, epochs)
+
+        # Extract accuracy from results
+        accuracy = model_results.get("accuracy")
+        val_results[name] = accuracy
+
+        # Extract and store metrics
+        if "roc_auc" in model_results and "micro" in model_results["roc_auc"]:
+            val_roc_auc_summary[name] = model_results["roc_auc"]["micro"]
+        else:
+            val_roc_auc_summary[name] = None
+
+        if "pr_auc" in model_results and "micro" in model_results["pr_auc"]:
+            val_pr_auc_summary[name] = model_results["pr_auc"]["micro"]
+        else:
+            val_pr_auc_summary[name] = None
+
+        # Store kappa score
+        if "kappa" in model_results:
+            val_kappa_summary[name] = model_results["kappa"]
+        else:
+            val_kappa_summary[name] = None
+
+        # Evaluate on test set
+        if accuracy is not None:
+            print(f"\n{'='*20} Testing {name} on Test Set {'='*20}\n")
+            test_model_results = evaluator.evaluate_model(model, test_loader)
+
+            # Extract accuracy from test results
+            test_accuracy = test_model_results.get("accuracy")
+            test_results[name] = test_accuracy
+
+            # Extract and store test metrics
+            if (
+                "roc_auc" in test_model_results
+                and "micro" in test_model_results["roc_auc"]
+            ):
+                test_roc_auc_summary[name] = test_model_results["roc_auc"]["micro"]
+            else:
+                test_roc_auc_summary[name] = None
+
+            if (
+                "pr_auc" in test_model_results
+                and "micro" in test_model_results["pr_auc"]
+            ):
+                test_pr_auc_summary[name] = test_model_results["pr_auc"]["micro"]
+            else:
+                test_pr_auc_summary[name] = None
+
+            # Store test kappa score
+            if "kappa" in test_model_results:
+                test_kappa_summary[name] = test_model_results["kappa"]
+            else:
+                test_kappa_summary[name] = None
+
+            # Track best model
+            if test_accuracy > best_accuracy:
+                best_accuracy = test_accuracy
+                best_model_obj = copy.deepcopy(model)
+                best_model_name = name
+
+    # Print comprehensive comparison
+    print("\n\n" + "=" * 100)
+    print("COMPREHENSIVE MODEL COMPARISON")
+    print("=" * 100)
+    print(
+        f"{'Model':<20}{'Val Acc':<10}{'Test Acc':<10}{'Val ROC AUC':<14}{'Test ROC AUC':<14}{'Val PR AUC':<14}{'Test PR AUC':<14}{'Val Kappa':<14}{'Test Kappa':<14}"
+    )
+    print("-" * 100)
+
+    for name in val_results.keys():
+        val_acc = val_results[name]
+        test_acc = test_results.get(name, None)
+        val_roc = val_roc_auc_summary.get(name, None)
+        test_roc = test_roc_auc_summary.get(name, None)
+        val_pr = val_pr_auc_summary.get(name, None)
+        test_pr = test_pr_auc_summary.get(name, None)
+        val_kappa = val_kappa_summary.get(name, None)
+        test_kappa = test_kappa_summary.get(name, None)
+
+        # Format values for display
+        val_acc_str = f"{val_acc:.4f}" if val_acc is not None else "Failed"
+        test_acc_str = f"{test_acc:.4f}" if test_acc is not None else "N/A"
+        val_roc_str = f"{val_roc:.4f}" if val_roc is not None else "N/A"
+        test_roc_str = f"{test_roc:.4f}" if test_roc is not None else "N/A"
+        val_pr_str = f"{val_pr:.4f}" if val_pr is not None else "N/A"
+        test_pr_str = f"{test_pr:.4f}" if test_pr is not None else "N/A"
+        val_kappa_str = f"{val_kappa:.4f}" if val_kappa is not None else "N/A"
+        test_kappa_str = f"{test_kappa:.4f}" if test_kappa is not None else "N/A"
+
+        print(
+            f"{name:<20}{val_acc_str:<10}{test_acc_str:<10}{val_roc_str:<14}{test_roc_str:<14}{val_pr_str:<14}{test_pr_str:<14}{val_kappa_str:<14}{test_kappa_str:<14}"
+        )
+
+    # Identify best model based on test metrics
+    if test_results:
+        # Best model by accuracy
+        best_acc_model = max(
+            test_results.items(), key=lambda x: x[1] if x[1] is not None else -1
+        )
+        print(
+            f"\nBest model by accuracy: {best_acc_model[0]} (Test Accuracy: {best_acc_model[1]:.4f})"
+        )
+
+        # Best model by ROC AUC (if available)
+        if any(v is not None for v in test_roc_auc_summary.values()):
+            best_roc_model = max(
+                [(k, v) for k, v in test_roc_auc_summary.items() if v is not None],
+                key=lambda x: x[1] if x[1] is not None else -1,
+            )
+            print(
+                f"Best model by ROC AUC: {best_roc_model[0]} (Test ROC AUC: {best_roc_model[1]:.4f})"
+            )
+
+        # Best model by PR AUC (if available)
+        if any(v is not None for v in test_pr_auc_summary.values()):
+            best_pr_model = max(
+                [(k, v) for k, v in test_pr_auc_summary.items() if v is not None],
+                key=lambda x: x[1] if x[1] is not None else -1,
+            )
+            print(
+                f"Best model by PR AUC: {best_pr_model[0]} (Test PR AUC: {best_pr_model[1]:.4f})"
+            )
+
+        # Best model by Kappa (if available)
+        if any(v is not None for v in test_kappa_summary.values()):
+            best_kappa_model = max(
+                [(k, v) for k, v in test_kappa_summary.items() if v is not None],
+                key=lambda x: x[1] if x[1] is not None else -1,
+            )
+            print(
+                f"Best model by Cohen's Kappa: {best_kappa_model[0]} (Test Kappa: {best_kappa_model[1]:.4f})"
+            )
+
+        # Save the best model (by accuracy)
+        if best_model_obj is not None:
+            try:
+                model_save_path = (
+                    f"best_model_{best_model_name.lower().replace(' ', '_')}.pth"
+                )
+                torch.save(best_model_obj.state_dict(), model_save_path)
+                print(f"Best model saved to {model_save_path}")
+            except Exception as save_error:
+                print(f"Error saving best model: {save_error}")
+    else:
+        print("\nNo models successfully completed testing.")
+
+    print("=" * 100)

utils/DatasetHandler.py

@@ -0,0 +1,93 @@
+import os
+import time
+import random
+import copy
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import torchvision.transforms as transforms
+from torchvision import transforms, datasets
+import torchvision.models as models
+import timm
+import matplotlib.pyplot as plt
+import seaborn as sns
+from PIL import Image
+from torch.utils.data import DataLoader, random_split, Subset, Dataset
+from sklearn.metrics import (
+    accuracy_score,
+    confusion_matrix,
+    classification_report,
+    roc_curve,
+    auc,
+    precision_recall_curve,
+    average_precision_score,
+)
+from sklearn.preprocessing import label_binarize
+from tqdm import tqdm
+import gc
+
+
+# Modified FilteredImageDataset class with Pterygium filtering
+class FilteredImageDataset(Dataset):
+    def __init__(self, dataset, excluded_classes=None):
+        """
+        Create a filtered dataset that excludes specific classes.
+
+        Args:
+            dataset: Original dataset (ImageFolder or similar)
+            excluded_classes: List of class names to exclude (e.g., ["Pterygium"])
+        """
+        self.dataset = dataset
+        self.excluded_classes = excluded_classes or []
+
+        # Get original class information
+        self.orig_classes = dataset.classes
+        self.orig_class_to_idx = dataset.class_to_idx
+
+        # Create indices of samples to keep (excluding specified classes)
+        self.indices = []
+        for idx, (_, target) in enumerate(dataset.samples):
+            class_name = self.orig_classes[target]
+            if class_name not in self.excluded_classes:
+                self.indices.append(idx)
+
+        # Create new class mapping without excluded classes
+        remaining_classes = [
+            c for c in self.orig_classes if c not in self.excluded_classes
+        ]
+        self.classes = remaining_classes
+        self.class_to_idx = {cls: idx for idx, cls in enumerate(remaining_classes)}
+        self.idx_to_class = {v: k for k, v in self.class_to_idx.items()}
+
+        # Create a mapping from old indices to new indices
+        self.target_mapping = {}
+        for old_class, old_idx in self.orig_class_to_idx.items():
+            if old_class in self.class_to_idx:
+                self.target_mapping[old_idx] = self.class_to_idx[old_class]
+
+        print(f"Filtered out classes: {self.excluded_classes}")
+        print(f"Remaining classes: {self.classes}")
+        print(
+            f"Original dataset size: {len(dataset)}, Filtered dataset size: {len(self.indices)}"
+        )
+
+    def __getitem__(self, index):
+        """Get item from the filtered dataset with remapped class labels."""
+        orig_idx = self.indices[index]
+        img, old_target = self.dataset[orig_idx]
+
+        # Remap target to new class index
+        new_target = self.target_mapping[old_target]
+
+        return img, new_target
+
+    def __len__(self):
+        """Return the number of samples in the filtered dataset."""
+        return len(self.indices)
+
+    # Allow transform to be updated
+    def set_transform(self, transform):
+        """Update the transform for the dataset."""
+        self.dataset.transform = transform

utils/Evaluator.py

@@ -0,0 +1,364 @@
+import numpy as np
+import torch
+import matplotlib.pyplot as plt
+import seaborn as sns
+from sklearn.metrics import (
+    accuracy_score,
+    classification_report,
+    confusion_matrix,
+    roc_curve,
+    precision_recall_curve,
+    auc,
+    average_precision_score,
+    cohen_kappa_score,
+)
+from sklearn.preprocessing import label_binarize
+
+
+class ClassificationEvaluator:
+    """
+    A class to evaluate and visualize classification model performance.
+
+    This class provides methods to compute various classification metrics
+    and generate visualizations for model evaluation.
+    """
+
+    def __init__(self, class_names):
+        """
+        Initialize the evaluator with class names.
+
+        Parameters:
+        - class_names: list of class names
+        """
+        self.class_names = class_names
+        self.num_classes = len(class_names)
+
+    def _ensure_numpy(self, data):
+        """Convert tensor to numpy if needed."""
+        if torch.is_tensor(data):
+            return data.cpu().numpy()
+        return np.array(data)
+
+    def evaluate_model(self, model, test_loader):
+        """
+        Evaluate a trained model on test dataset.
+
+        Parameters:
+        - model: PyTorch model to evaluate
+        - test_loader: DataLoader containing test data
+
+        Returns:
+        - results: Dictionary containing evaluation metrics
+        """
+        model.eval()
+        device = next(model.parameters()).device
+
+        all_labels = []
+        all_preds = []
+        all_scores = []
+
+        with torch.no_grad():
+            for inputs, labels in test_loader:
+                inputs, labels = inputs.to(device), labels.to(device)
+                outputs = model(inputs)
+                _, preds = torch.max(outputs, 1)
+
+                all_labels.extend(labels.cpu().numpy())
+                all_preds.extend(preds.cpu().numpy())
+                all_scores.append(
+                    torch.nn.functional.softmax(outputs, dim=1).cpu().numpy()
+                )
+
+        all_scores = np.vstack(all_scores)
+
+        # Compute metrics
+        results = self.compute_metrics(all_labels, all_preds, all_scores)
+        return results
+
+    def compute_metrics(self, y_true, y_pred, y_scores, model_name=""):
+        """
+        Compute comprehensive classification metrics.
+
+        Parameters:
+        - y_true: true labels
+        - y_pred: predicted labels
+        - y_scores: predicted probability scores
+        - model_name: name of the model (optional)
+
+        Returns:
+        - Dictionary containing all metrics
+        """
+        # Ensure numpy arrays
+        y_true = self._ensure_numpy(y_true)
+        y_pred = self._ensure_numpy(y_pred)
+        y_scores = self._ensure_numpy(y_scores)
+
+        # Calculate accuracy
+        accuracy = accuracy_score(y_true, y_pred)
+        print(f"Overall Accuracy: {accuracy:.4f}")
+
+        # Calculate and display Cohen's Kappa
+        kappa = cohen_kappa_score(y_true, y_pred)
+        print(f"Cohen's Kappa Score: {kappa:.4f}")
+
+        # Generate classification report
+        report = classification_report(
+            y_true, y_pred, target_names=self.class_names, output_dict=True
+        )
+
+        # Print formatted classification report
+        print("\nClassification Report:")
+        print(classification_report(y_true, y_pred, target_names=self.class_names))
+
+        # Calculate ROC curves and AUC for each class
+        print("\nCalculating ROC curves...")
+        roc_auc_dict = self.plot_roc_curves(y_true, y_scores)
+
+        # Calculate PR curves and AUC for each class
+        print("\nCalculating PR curves...")
+        pr_auc_dict = self.plot_pr_curves(y_true, y_scores)
+
+        # Plot confusion matrix
+        print("\nGenerating confusion matrix...")
+        self.plot_confusion_matrix(y_true, y_pred)
+
+        # Plot per-class accuracy
+        print("\nCalculating per-class accuracy...")
+        self.plot_per_class_accuracy(y_true, y_pred)
+
+        # Return metrics dictionary
+        return {
+            "accuracy": accuracy,
+            "report": report,
+            "roc_auc": roc_auc_dict,
+            "pr_auc": pr_auc_dict,
+            "kappa": kappa,
+        }
+
+    def plot_roc_curves(self, y_true, y_scores):
+        """
+        Plot ROC curves for multi-class classification.
+
+        Parameters:
+        - y_true: true labels
+        - y_scores: predicted probability scores
+
+        Returns:
+        - Dictionary containing AUC values for each class
+        """
+        y_true = self._ensure_numpy(y_true)
+        y_scores = self._ensure_numpy(y_scores)
+
+        # Binarize the labels for one-vs-rest ROC calculation
+        y_true_bin = label_binarize(y_true, classes=range(self.num_classes))
+
+        # Compute ROC curve and ROC area for each class
+        fpr = {}
+        tpr = {}
+        roc_auc = {}
+
+        plt.figure(figsize=(12, 8))
+
+        for i in range(self.num_classes):
+            fpr[i], tpr[i], _ = roc_curve(y_true_bin[:, i], y_scores[:, i])
+            roc_auc[i] = auc(fpr[i], tpr[i])
+
+            plt.plot(
+                fpr[i],
+                tpr[i],
+                lw=2,
+                label=f"{self.class_names[i]} (area = {roc_auc[i]:.2f})",
+            )
+
+        # Plot the diagonal (random classifier)
+        plt.plot([0, 1], [0, 1], "k--", lw=2)
+
+        # Calculate and plot micro-average ROC curve
+        fpr["micro"], tpr["micro"], _ = roc_curve(y_true_bin.ravel(), y_scores.ravel())
+        roc_auc["micro"] = auc(fpr["micro"], tpr["micro"])
+        plt.plot(
+            fpr["micro"],
+            tpr["micro"],
+            label=f'Micro-average (area = {roc_auc["micro"]:.2f})',
+            lw=2,
+            linestyle=":",
+            color="deeppink",
+        )
+
+        plt.xlim([0.0, 1.0])
+        plt.ylim([0.0, 1.05])
+        plt.xlabel("False Positive Rate")
+        plt.ylabel("True Positive Rate")
+        plt.title("ROC Curves")
+        plt.legend(loc="lower right")
+        plt.grid(True, alpha=0.3)
+        plt.tight_layout()
+        plt.show()
+
+        return roc_auc
+
+    def plot_pr_curves(self, y_true, y_scores):
+        """
+        Plot Precision-Recall curves for multi-class classification.
+
+        Parameters:
+        - y_true: true labels
+        - y_scores: predicted probability scores
+
+        Returns:
+        - Dictionary containing average precision values for each class
+        """
+        y_true = self._ensure_numpy(y_true)
+        y_scores = self._ensure_numpy(y_scores)
+
+        # Binarize the labels
+        y_true_bin = label_binarize(y_true, classes=range(self.num_classes))
+
+        # Compute PR curve and average precision for each class
+        precision = {}
+        recall = {}
+        avg_precision = {}
+
+        plt.figure(figsize=(12, 8))
+
+        for i in range(self.num_classes):
+            precision[i], recall[i], _ = precision_recall_curve(
+                y_true_bin[:, i], y_scores[:, i]
+            )
+            avg_precision[i] = average_precision_score(y_true_bin[:, i], y_scores[:, i])
+
+            plt.plot(
+                recall[i],
+                precision[i],
+                lw=2,
+                label=f"{self.class_names[i]} (AP = {avg_precision[i]:.2f})",
+            )
+
+        # Calculate and plot micro-average PR curve
+        precision["micro"], recall["micro"], _ = precision_recall_curve(
+            y_true_bin.ravel(), y_scores.ravel()
+        )
+        avg_precision["micro"] = average_precision_score(
+            y_true_bin.ravel(), y_scores.ravel()
+        )
+
+        plt.plot(
+            recall["micro"],
+            precision["micro"],
+            label=f'Micro-average (AP = {avg_precision["micro"]:.2f})',
+            lw=2,
+            linestyle=":",
+            color="deeppink",
+        )
+
+        plt.xlim([0.0, 1.0])
+        plt.ylim([0.0, 1.05])
+        plt.xlabel("Recall")
+        plt.ylabel("Precision")
+        plt.title("Precision-Recall Curves")
+        plt.legend(loc="best")
+        plt.grid(True, alpha=0.3)
+        plt.tight_layout()
+        plt.show()
+
+        return avg_precision
+
+    def plot_confusion_matrix(self, y_true, y_pred):
+        """
+        Plot confusion matrix.
+
+        Parameters:
+        - y_true: true labels
+        - y_pred: predicted labels
+        """
+        y_true = self._ensure_numpy(y_true)
+        y_pred = self._ensure_numpy(y_pred)
+
+        # Get unique values in both arrays
+        unique_values = np.unique(np.concatenate([y_true, y_pred]))
+        print(f"Unique values in confusion matrix data: {unique_values}")
+
+        # Create the confusion matrix with explicit labels
+        cm = confusion_matrix(y_true, y_pred, labels=range(self.num_classes))
+
+        plt.figure(figsize=(10, 8))
+        sns.heatmap(
+            cm,
+            annot=True,
+            fmt="d",
+            cmap="Blues",
+            xticklabels=self.class_names,
+            yticklabels=self.class_names,
+        )
+        plt.title("Confusion Matrix")
+        plt.xlabel("Predicted")
+        plt.ylabel("True")
+        plt.tight_layout()
+        plt.show()
+
+    def plot_per_class_accuracy(self, y_true, y_pred):
+        """
+        Plot per-class accuracy.
+
+        Parameters:
+        - y_true: true labels
+        - y_pred: predicted labels
+        """
+        y_true = self._ensure_numpy(y_true)
+        y_pred = self._ensure_numpy(y_pred)
+
+        # Create the confusion matrix with explicit labels
+        cm = confusion_matrix(y_true, y_pred, labels=range(self.num_classes))
+
+        # Calculate per-class accuracy
+        per_class_accuracy = np.zeros(self.num_classes)
+        for i in range(self.num_classes):
+            if i < cm.shape[0] and np.sum(cm[i, :]) > 0:
+                per_class_accuracy[i] = cm[i, i] / np.sum(cm[i, :])
+
+        # Create the bar plot
+        plt.figure(figsize=(14, 7))
+        plt.bar(range(self.num_classes), per_class_accuracy, color="skyblue")
+        plt.xticks(range(self.num_classes), self.class_names, rotation=45, ha="right")
+        plt.xlabel("Classes")
+        plt.ylabel("Accuracy")
+        plt.title("Per-Class Accuracy")
+        plt.tight_layout()
+        plt.show()
+
+        return per_class_accuracy
+
+    def plot_training_history(self, train_losses, val_losses, train_accs, val_accs):
+        """
+        Plot accuracy and loss curves from training history.
+
+        Parameters:
+        - train_losses: list of training losses
+        - val_losses: list of validation losses
+        - train_accs: list of training accuracies
+        - val_accs: list of validation accuracies
+        """
+        plt.figure(figsize=(12, 5))
+
+        # Accuracy curve
+        plt.subplot(1, 2, 1)
+        plt.plot(train_accs, label="Train Accuracy")
+        plt.plot(val_accs, label="Validation Accuracy")
+        plt.xlabel("Epochs")
+        plt.ylabel("Accuracy")
+        plt.title("Accuracy Curve")
+        plt.legend()
+        plt.grid(True)
+
+        # Loss curve
+        plt.subplot(1, 2, 2)
+        plt.plot(train_losses, label="Train Loss")
+        plt.plot(val_losses, label="Validation Loss")
+        plt.xlabel("Epochs")
+        plt.ylabel("Loss")
+        plt.title("Loss Curve")
+        plt.legend()
+        plt.grid(True)
+
+        plt.tight_layout()
+        plt.show()

utils/ModelCreator.py

@@ -0,0 +1,180 @@
+import torch
+import torch.nn as nn
+import timm
+
+# Set device
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class EyeDetectionModels:
+    def __init__(self, num_classes, freeze_layers=True, device=DEVICE):
+        """
+        Initialize the EyeDetectionModels class.
+        This class provides methods to create and configure various deep learning models for eye detection.
+        """
+        # Initialize the model creator
+        self.num_classes = num_classes
+        self.freeze_layers = freeze_layers
+        self.device = device
+        self.models = {
+            "mobilenetv4": self.get_model_mobilenetv4,
+            "levit": self.get_model_levit,
+            "efficientvit": self.get_model_efficientvit,
+            "gernet": self.get_model_gernet,
+            "regnetx": self.get_model_regnetx,
+        }
+
+    # Model architecture functions
+    @staticmethod
+    def _get_feature_blocks(model):
+        """
+        Utility: locate the main feature blocks container in a timm model.
+        Returns a list-like module of blocks.
+        """
+        for attr in ("features", "blocks", "layers", "stem"):  # common container names
+            if hasattr(model, attr):
+                return getattr(model, attr)
+        # fallback: collect all children except classifier/head
+        return list(model.children())[:-1]
+
+    @staticmethod
+    def _freeze_except_last_n(blocks, n):
+        total = len(blocks)
+        for idx, block in enumerate(blocks):
+            requires = idx >= total - n
+            for p in block.parameters():
+                p.requires_grad = requires
+
+    def get_model_mobilenetv4(self):
+        model = timm.create_model(
+            "mobilenetv4_conv_medium.e500_r256_in1k", pretrained=True
+        )
+        if self.freeze_layers:
+            blocks = self._get_feature_blocks(model)
+            self._freeze_except_last_n(blocks, 2)
+        # replace classifier
+        in_features = model.classifier.in_features
+        model.classifier = nn.Sequential(
+            nn.Linear(in_features, 512),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.4),
+            nn.Linear(512, self.num_classes),
+        )
+        return model.to(self.device)
+
+    def get_model_levit(self):
+        model = timm.create_model("levit_128s.fb_dist_in1k", pretrained=True)
+        if self.freeze_layers:
+            blocks = self._get_feature_blocks(model)
+            self._freeze_except_last_n(blocks, 2)
+        # Attempt to extract in_features from model.head or classifier
+        head = getattr(model, "head_dist", None) or getattr(model, "classifier", None)
+        linear = getattr(head, "linear")
+        in_features = 384
+        model.head = nn.Sequential(
+            nn.Linear(in_features, 512),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.4),
+            nn.Linear(512, self.num_classes),
+        )
+        model.head_dist = nn.Sequential(
+            nn.Linear(in_features, 512),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.4),
+            nn.Linear(512, self.num_classes),
+        )
+        return model.to(self.device)
+
+    def get_model_efficientvit(self):
+        model = timm.create_model("efficientvit_m1.r224_in1k", pretrained=True)
+        if self.freeze_layers:
+            blocks = self._get_feature_blocks(model)
+            self._freeze_except_last_n(blocks, 2)
+        # handle different head naming
+        head = getattr(model, "head", None)
+        print(head)
+        linear = getattr(head, "linear")
+        in_features = 192
+        model.head.linear = nn.Sequential(
+            nn.Linear(in_features, 512),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.4),
+            nn.Linear(512, self.num_classes),
+        )
+        return model.to(self.device)
+
+    def get_model_gernet(self):
+        """
+        Load and configure a GENet (General and Efficient Network) model with customizable classifier.
+
+        Returns:
+            Configured GENet model
+        """
+        model = timm.create_model("gernet_s.idstcv_in1k", pretrained=True)
+
+        if self.freeze_layers:
+            # For GENet, we need to specifically handle its structure
+            # It typically has a 'stem' and 'stages' structure
+            if hasattr(model, "stem") and hasattr(model, "stages"):
+                # Freeze stem completely
+                for param in model.stem.parameters():
+                    param.requires_grad = False
+
+                # Freeze all stages except the last two
+                stages = list(model.stages.children())
+                total_stages = len(stages)
+                for i, stage in enumerate(stages):
+                    requires_grad = i >= total_stages - 2
+                    for param in stage.parameters():
+                        param.requires_grad = requires_grad
+            else:
+                # Fallback to generic approach
+                blocks = self._get_feature_blocks(model)
+                self._freeze_except_last_n(blocks, 2)
+
+        # Replace classifier
+        in_features = model.head.fc.in_features
+        model.head.fc = nn.Sequential(
+            nn.Linear(in_features, 512),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.4),
+            nn.Linear(512, self.num_classes),
+        )
+        return model.to(self.device)
+
+    def get_model_regnetx(self):
+        """
+        Load and configure a RegNetX model with customizable classifier.
+
+        Returns:
+            Configured RegNetX model
+        """
+        model = timm.create_model("regnetx_008.tv2_in1k", pretrained=True)
+
+        if self.freeze_layers:
+            for param in model.parameters():
+                param.requires_grad = False
+
+            # RegNetX typically has 'stem' + 'trunk' structure in timm
+            if hasattr(model, "trunk"):
+                # Unfreeze final stages of the trunk
+                trunk_blocks = list(model.trunk.children())
+                # Unfreeze approximately last 25% of trunk blocks
+                unfreeze_from = max(0, int(len(trunk_blocks) * 0.75))
+                for i in range(unfreeze_from, len(trunk_blocks)):
+                    for param in trunk_blocks[i].parameters():
+                        param.requires_grad = True
+
+            # Always unfreeze the classifier/head for fine-tuning
+            for param in model.head.parameters():
+                param.requires_grad = True
+
+        # Replace classifier
+        in_features = model.head.fc.in_features
+        model.head.fc = nn.Sequential(
+            nn.Linear(in_features, 512),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.4),
+            nn.Linear(512, self.num_classes),
+        )
+        return model.to(self.device)

utils/Trainer.py

@@ -0,0 +1,285 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torch.optim as optim
+import torch.nn as nn
+from tqdm import tqdm
+import gc
+
+from Evaluator import ClassificationEvaluator
+from Callback import EarlyStopping
+
+
+def train_model(
+    model,
+    criterion,
+    optimizer,
+    scheduler,
+    train_loader,
+    val_loader,
+    early_stopping,
+    epochs=15,
+    use_ddp=False,
+):
+    """
+    Train the model and perform validation using multiple GPUs.
+    Supports both DataParallel (DP) and DistributedDataParallel (DDP) modes.
+
+    Args:
+        model: Model to train
+        criterion: Loss function
+        optimizer: Optimizer for training
+        scheduler: Learning rate scheduler
+        train_loader: DataLoader for training data
+        val_loader: DataLoader for validation data
+        early_stopping: Early stopping handler
+        epochs: Maximum number of epochs to train
+        use_ddp: Whether to use DistributedDataParallel (True) or DataParallel (False)
+    """
+    # Check available GPUs
+    num_gpus = torch.cuda.device_count()
+    if num_gpus < 2:
+        print(
+            f"Warning: Requested multi-GPU training but only {num_gpus} GPU(s) available. Continuing with available resources."
+        )
+    else:
+        print(f"Using {num_gpus} GPUs for training")
+
+    # Setup device and model
+    if num_gpus >= 2:
+        if use_ddp:
+            # For DistributedDataParallel
+            import torch.distributed as dist
+            from torch.nn.parallel import DistributedDataParallel as DDP
+
+            # Initialize process group
+            dist.init_process_group(backend="nccl")
+            local_rank = dist.get_rank()
+            torch.cuda.set_device(local_rank)
+            device = torch.device(f"cuda:{local_rank}")
+
+            model = model.to(device)
+            model = DDP(model, device_ids=[local_rank])
+        else:
+            # For DataParallel (simpler to use)
+            device = torch.device("cuda:0")
+            model = model.to(device)
+            model = torch.nn.DataParallel(model)
+    else:
+        # Single GPU
+        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        model = model.to(device)
+
+    train_losses = []
+    val_losses = []
+    train_accs = []
+    val_accs = []
+
+    # Store validation predictions and labels for final evaluation
+    all_val_labels = []
+    all_val_preds = []
+    all_val_scores = []
+
+    for epoch in range(epochs):
+        print(f"Epoch {epoch+1}/{epochs}")
+
+        # Training phase
+        model.train()
+        running_loss = 0.0
+        correct = 0
+        total = 0
+
+        for inputs, labels in tqdm(train_loader, desc="Training"):
+            inputs, labels = inputs.to(device), labels.to(device)
+
+            optimizer.zero_grad()
+            outputs = model(inputs)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+
+            running_loss += loss.item() * inputs.size(0)
+            _, predicted = torch.max(outputs, 1)
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+
+        epoch_train_loss = running_loss / len(train_loader.dataset)
+        epoch_train_acc = correct / total
+        train_losses.append(epoch_train_loss)
+        train_accs.append(epoch_train_acc)
+
+        # Validation phase
+        model.eval()
+        running_loss = 0.0
+        correct = 0
+        total = 0
+
+        all_labels = []
+        all_preds = []
+        all_scores = []
+
+        with torch.no_grad():
+            for inputs, labels in tqdm(val_loader, desc="Validation"):
+                inputs, labels = inputs.to(device), labels.to(device)
+                outputs = model(inputs)
+                loss = criterion(outputs, labels)
+
+                running_loss += loss.item() * inputs.size(0)
+                probs = F.softmax(outputs, dim=1)
+                _, predicted = torch.max(outputs, 1)
+                total += labels.size(0)
+                correct += (predicted == labels).sum().item()
+
+                all_labels.extend(labels.cpu().numpy().tolist())
+                all_preds.extend(predicted.cpu().numpy().tolist())
+                all_scores.append(probs.cpu().numpy())
+
+        epoch_val_loss = running_loss / len(val_loader.dataset)
+        epoch_val_acc = correct / total
+        val_losses.append(epoch_val_loss)
+        val_accs.append(epoch_val_acc)
+
+        all_scores = np.vstack(all_scores) if all_scores else np.array([])
+
+        # Store validation results for the final epoch
+        all_val_labels = all_labels
+        all_val_preds = all_preds
+        all_val_scores = all_scores
+
+        # Update learning rate scheduler
+        scheduler.step(epoch_val_loss)
+
+        print(f"Train Loss: {epoch_train_loss:.4f}, Train Acc: {epoch_train_acc:.4f}")
+        print(f"Val Loss: {epoch_val_loss:.4f}, Val Acc: {epoch_val_acc:.4f}")
+        print(f"Learning rate: {optimizer.param_groups[0]['lr']:.6f}")
+
+        # Check early stopping
+        early_stopping(epoch_val_loss)
+        if early_stopping.early_stop:
+            print("Early stopping triggered!")
+            break
+
+        # Free up memory
+        del all_labels, all_preds, all_scores
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    # Clean up DDP if used
+    if num_gpus >= 2 and use_ddp:
+        dist.destroy_process_group()
+
+    return (
+        model,
+        train_losses,
+        val_losses,
+        train_accs,
+        val_accs,
+        all_val_labels,
+        all_val_preds,
+        all_val_scores,
+    )
+
+
+def model_train(model, train_loader, val_loader, dataset, epochs=20):
+    model_name = type(model).__name__
+    if hasattr(model, "pretrained_cfg") and "name" in model.pretrained_cfg:
+        model_name = model.pretrained_cfg["name"]
+
+    print(f"\n{'='*20} Training {model_name} {'='*20}\n")
+
+    class_names = dataset.classes
+    num_classes = len(class_names)
+    learning_rate = 0.001
+
+    try:
+        optimizer = optim.Adam(model.parameters(), lr=learning_rate)
+        scheduler = optim.lr_scheduler.ReduceLROnPlateau(
+            optimizer, mode="min", factor=0.1, patience=3
+        )
+        early_stopping = EarlyStopping(patience=5)
+
+        (
+            model,
+            train_losses,
+            val_losses,
+            train_accs,
+            val_accs,
+            val_labels,
+            val_preds,
+            val_scores,
+        ) = train_model(
+            model,
+            nn.CrossEntropyLoss(),
+            optimizer,
+            scheduler,
+            train_loader,
+            val_loader,
+            early_stopping,
+            epochs=epochs,
+            use_ddp=False,
+        )
+
+        print(f"\n{'='*20} Evaluation for {model_name} {'='*20}\n")
+        evaluator = ClassificationEvaluator(
+            num_classes=num_classes,
+            class_names=class_names,
+        )
+
+        evaluator.plot_training_history(train_losses, val_losses, train_accs, val_accs)
+        # Process validation predictions and labels
+        try:
+            evaluator.plot_confusion_matrix(val_labels, val_preds)
+            evaluator.plot_per_class_accuracy(val_labels, val_preds)
+
+            # Get metrics from the updated function including kappa
+            accuracy, report_dict, roc_auc_dict, pr_auc_dict, kappa = (
+                evaluator.compute_metrics(
+                    val_labels,
+                    val_preds,
+                    val_scores,
+                    model_name,
+                )
+            )
+
+            # Build a results dictionary including kappa
+            results = {
+                "accuracy": accuracy,
+                "report": report_dict,
+                "roc_auc": roc_auc_dict,
+                "pr_auc": pr_auc_dict,
+                "kappa": kappa,
+            }
+
+            return results
+        except Exception as viz_error:
+            print(f"Error in visualization: {viz_error}")
+            import traceback
+
+            traceback.print_exc()
+            return {"accuracy": None}
+
+    except Exception as e:
+        print(f"Error occurred when training {model_name}: {e}")
+        import traceback
+
+        traceback.print_exc()
+        return {"accuracy": None}
+    finally:
+        # Clean up memory
+        if "optimizer" in locals():
+            del optimizer
+        if "scheduler" in locals():
+            del scheduler
+        if "early_stopping" in locals():
+            del early_stopping
+        if "train_losses" in locals():
+            del train_losses
+            del val_losses
+            del train_accs
+            del val_accs
+            del val_labels
+            del val_preds
+            del val_scores
+
+        gc.collect()
+        torch.cuda.empty_cache()