Uploaded

HAM_10000_CLASSIFICATION-master/.devcontainer/Dockerfile

@@ -0,0 +1,17 @@
+FROM python:3.12-bookworm
+
+# Install system dependencies (libGL, FFmpeg, X11 tools)
+RUN apt-get update && \
+    apt-get install -y libgl1 ffmpeg x11-xserver-utils && \
+    apt-get clean && rm -rf /var/lib/apt/lists/*
+
+# Install Python dependencies (optimized with --no-cache-dir)
+RUN pip install --no-cache-dir \
+    numpy \
+    pandas \
+    matplotlib \
+    opencv-python-headless \
+    torch \
+    torchvision \
+    streamlit \
+    tqdm

HAM_10000_CLASSIFICATION-master/.devcontainer/devcontainer.json

@@ -0,0 +1,47 @@
+// For format details, see https://aka.ms/devcontainer.json. For config options, see the
+// README at: https://github.com/devcontainers/templates/tree/main/src/python
+{
+	"name": "Python 3",
+	// Or use a Dockerfile or Docker Compose file. More info: https://containers.dev/guide/dockerfile
+	
+	"build" : {
+		"dockerfile": "Dockerfile"
+	},
+
+	// Features to add to the dev container. More info: https://containers.dev/features.
+	// "features": {},
+	"privileged" : true,
+
+	// Use 'forwardPorts' to make a list of ports inside the container available locally.
+	"forwardPorts": [
+		8501
+	],
+
+	"runArgs": [
+		"--env=DISPLAY",
+		"--env=XAUTHORITY",
+		"--env=WAYLAND_DISPLAY",   // Added Wayland environment variable
+		"--env=QT_QPA_PLATFORM",   // Ensures Qt apps use the right display server
+		"--volume=/tmp/.X11-unix:/tmp/.X11-unix",
+		"--volume=/run/user/1000/.mutter-Xwaylandauth.HSC612",
+		"--device=/dev/dri:/dev/dri"
+	],
+
+	"customizations": {
+		"vscode": {
+			"extensions": [
+				"ms-python.python",
+				"ms-python.debugpy"
+			]
+		}
+	}
+
+	// Use 'postCreateCommand' to run commands after the container is created.
+	// "postCreateCommand": "pip3 install --user -r requirements.txt",
+
+	// Configure tool-specific properties.
+	// "customizations": {},
+
+	// Uncomment to connect as root instead. More info: https://aka.ms/dev-containers-non-root.
+	// "remoteUser": "root"
+}

HAM_10000_CLASSIFICATION-master/.dockerignore

@@ -0,0 +1,7 @@
+.devcontainer
+dataset
+__pycache__
+utils
+venv
+test.py
+README.md

HAM_10000_CLASSIFICATION-master/.gitignore

@@ -0,0 +1,8 @@
+dataset/
+model_weights/
+__pycache__/
+*.csv
+*.zip
+test.py
+venv/
+*.jpg

HAM_10000_CLASSIFICATION-master/Dockerfile

@@ -0,0 +1,17 @@
+FROM python:3.12-bookworm
+
+# Set the working directory
+WORKDIR /app
+
+# Copy the current directory contents into the container at /app
+COPY . /app
+
+# update ubuntu and pip
+RUN python -m pip install --upgrade pip
+
+RUN pip --default-timeout=100 --no-cache-dir install -r requirements.txt
+
+EXPOSE 8501
+
+CMD ["python", "app.py"]
+

HAM_10000_CLASSIFICATION-master/README.md

@@ -0,0 +1,128 @@
+# Ham10000-Backend
+
+This repo contains the Backend.
+
+
+
+## Installation
+
+- create a virtual environment.
+```bash
+$ python -m venv venv
+```
+
+- Install the required packages
+```bash
+$ pip install requirements.txt
+```
+
+- Run the backend.
+```bash
+$ python app.py
+```
+
+- if you see output like this. it means backend is running
+```bash
+INFO:     Started server process [1624]
+INFO:     Waiting for application startup.
+INFO:     Application startup complete.
+INFO:     Uvicorn running on http://0.0.0.0:8501 (Press CTRL+C to quit)
+```
+
+## API EndPoints
+
+- ### **GET** -  http://localhost:8501/api/model_performance
+
+  This endpoint is for details of each model. 
+
+  ```bash
+  {
+        "models": [
+            {
+                "name": "Efficient Net V7",
+                "description": "A high-efficiency architecture that leverages compound scaling for superior performance across various tasks.",
+                "performance_tags": [
+                    { "icon": "fa-tachometer-alt", "label": "High Efficiency" },
+                    { "icon": "fa-star", "label": "State-of-the-Art" }
+                ]
+            }
+        ]
+    }
+  ```
+
+  - ### POST  -  http://localhost:8501/api/classify
+    This endpoint is for making posting images to models and making predictions.
+
+    **Request Body**
+
+    - Json: 
+    ```bash
+    {
+      'image': 'binary file (JPEG, PNG, etc.)',
+      'model': 'Efficient Net V7'
+    }
+    ```
+
+    - CURL:
+    ```bash
+    curl -X POST "http://your-api-domain.com/api/classify" \
+     -H "Content-Type: multipart/form-data" \
+     -F "image=@path/to/image.jpg" \
+     -F "model=Efficient Net V7"
+    ```
+
+    **Response**
+
+    - If no cancer detected
+
+    ```bash
+      {
+        "result" : "Cancer Not Detected"
+      }
+    ```
+
+    - If cancer detected. The class with highest confidence is the prediction. In this case **akiec** has the highest confidence of **0.9999591112136841**.
+
+    ```bash
+    {
+      "result": [
+        {
+          "class": "akiec",
+          "confidence": 0.9999591112136841
+        },
+        {
+          "class": "mel",
+          "confidence": 0.00003423781890887767
+        },
+        {
+          "class": "bkl",
+          "confidence": 0.0000065627805270196404
+        },
+        {
+          "class": "nv",
+          "confidence": 1.223657761784125e-7
+        },
+        {
+          "class": "vasc",
+          "confidence": 7.100419363581523e-9
+        },
+        {
+          "class": "df",
+          "confidence": 6.403973351609693e-9
+        },
+        {
+          "class": "bcc",
+          "confidence": 1.5403813780068276e-9
+        }
+      ]
+    }
+    ```
+  
+  - ### GET  -  http://localhost:8501/api/health
+    This endpoint is for aws healthcheck only and doesn't have any partical significance.
+
+## Note:
+
+**'Efficient Net V7'** is the best model so use that for classification.  
+
+

HAM_10000_CLASSIFICATION-master/app.py

@@ -0,0 +1,169 @@
+import io
+from PIL import Image
+from models.efficient_net import EfficientNetB7
+from fastapi import FastAPI, File, UploadFile, Form, HTTPException
+from fastapi.responses import JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from contextlib import asynccontextmanager
+from torchvision import models, transforms
+import torch.nn.functional as F
+import torch
+import gdown
+import os
+from colorama import Fore, Style
+
+####################### DOWNLOAD MODEL WEIGHTS #######################
+if not os.path.exists("model_weights"):
+    drive_link = "https://drive.google.com/drive/folders/1JOd3O1c3me5JWE3Elq0MhhShP_XhTYRV"
+    output_dir = "model_weights"
+
+    # Download the folder recursively
+    gdown.download_folder(drive_link, output=output_dir, quiet=False, use_cookies=False)
+
+    # Check if model weights were downloaded
+    model_weights_path = os.path.join(output_dir, "model_weights/EfficientNetV7Large_v1/saved_models/best_test_model.pth")  # Update with actual filename
+
+    if os.path.exists(model_weights_path):
+        print(Fore.GREEN + "✅ Model weights downloaded successfully!" + Style.RESET_ALL)
+    else:
+        print(Fore.RED + "❌ Model weights not found. Check the folder or link." + Style.RESET_ALL)
+
+    print("Download completed!")
+else:
+    print(Fore.YELLOW + "⚠️ Model weights already exist. Skipping download." + Style.RESET_ALL)
+    print("Download skipped!")
+######################################################################
+
+ml_models = {}
+cancer_nocancer_model_1 = None
+cancer_nocancer_model_2 = None
+
+################### for cancer no cancer model ################################
+def load_model(model_path, num_classes):
+    """
+    Load the EfficientNet-B0 model from the state dict
+    """
+    model = models.efficientnet_b0(weights=models.EfficientNet_B0_Weights.IMAGENET1K_V1)
+    model.classifier[1] = torch.nn.Linear(model.classifier[1].in_features, num_classes)
+
+    state_dict = torch.load(model_path, map_location="cpu")
+    model.load_state_dict(state_dict)
+    model.eval()
+
+    return model
+
+efficient_net_b0_transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                            std=[0.229, 0.224, 0.225])
+])
+cancer_no_caner_class_mapping = {0: 'cancer', 1: 'nocancer'}
+
+################################################################################
+
+
+# make a lifespan
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    global ml_models
+    global cancer_nocancer_model_1
+    global cancer_nocancer_model_2
+
+    cancer_nocancer_1_weight_path = "model_weights/cancer_nocancer.pth" 
+    cancer_nocancer_2_weight_path = "model_weights/cancer_nocancer_100.pth"
+
+    efficient_net_model = EfficientNetB7(weights_path="model_weights/EfficientNetV7Large_v1/saved_models/best_test_model.pth")
+
+    cancer_nocancer_model_1 = load_model(cancer_nocancer_1_weight_path, len(cancer_no_caner_class_mapping))
+    cancer_nocancer_model_2 = load_model(cancer_nocancer_2_weight_path, len(cancer_no_caner_class_mapping))
+
+    ml_models = {
+        "EFFICIENT NET V7": efficient_net_model,
+    }
+    yield
+
+    ml_models.clear()
+
+app = FastAPI(lifespan=lifespan)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Allow all origins (or specify your frontend URL)
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+@app.post("/api/classify", response_class=JSONResponse)
+async def classify(
+    image: UploadFile = File(...),  # Accepts an image file
+    model: str = Form(...)         # Accepts a model name as form data
+):
+    global ml_models
+    global cancer_nocancer_model_1
+    global cancer_nocancer_model_2
+
+    model = model.upper()
+    if model not in ml_models.keys():
+        raise HTTPException(status_code=400, detail="Invalid model specified")
+
+    image_data = await image.read()
+    img = Image.open(io.BytesIO(image_data)).convert("RGB")
+
+    # check whether image if of cancer or not
+    img_for_check = efficient_net_b0_transform(img).unsqueeze(0)
+    check_cancer_1 = F.softmax(cancer_nocancer_model_1(img_for_check), dim=1)
+    check_cancer_2 = F.softmax(cancer_nocancer_model_2(img_for_check), dim=1)
+
+    check_cancer_1_index = check_cancer_1.argmax()
+    check_cancer_2_index = check_cancer_2.argmax()
+
+    print("-" * 25, "Cancer No Cancer", "-" * 25)
+    print(check_cancer_1_index)
+    print(check_cancer_2_index)
+    print("-" * 50) 
+
+    if check_cancer_1_index == 1 and check_cancer_2_index == 1:
+        print("Cancer Not Detected")
+        return JSONResponse({"result": "Cancer Not Detected"})
+    else:
+        print("Cancer Detected")
+        prediction = predict(ml_models[model], model, img)
+        return JSONResponse(prediction)
+    
+
+def predict(model, model_name, image_data):
+    classes = ['akiec','bcc', 'bkl', 'df', 'mel', 'nv', 'vasc']
+
+    print("-" * 50)
+    print(f"Model Used: {model_name}")
+    pred = model.make_prediction(image_data)[0]
+    print(pred)
+    pred_json = {"result" : [{"class": cls, "confidence": float(pred[i])} for i, cls in enumerate(classes)]}
+    pred_json["result"] = sorted(pred_json["result"], key=lambda x: x["confidence"], reverse=True)
+    print(pred_json)
+    print("-" * 50)
+    return pred_json
+
+
+@app.get("/api/model_performance", response_class=JSONResponse)
+async def home():
+    return {
+        "models": [
+            {
+                "name": "Efficient Net V7",
+                "description": "A high-efficiency architecture that leverages compound scaling for superior performance across various tasks.",
+                "performance_tags": [
+                    { "icon": "fa-tachometer-alt", "label": "High Efficiency" },
+                    { "icon": "fa-star", "label": "State-of-the-Art" }
+                ]
+            },
+        ]
+    }
+
+@app.get("/api/health", response_class=JSONResponse)
+def health():
+    return JSONResponse(status_code=200, content={"status": "Working fine!"})
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8501)

HAM_10000_CLASSIFICATION-master/models/convnext.py

@@ -0,0 +1,25 @@
+import torch
+from torchvision.models import convnext_base, ConvNeXt_Base_Weights
+from PIL import Image
+import torch.nn.functional as F
+
+class ConvNextBase:
+    def __init__(self, weights_path: str):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        self.model = convnext_base()
+        self.model.classifier[2] = torch.nn.Linear(self.model.classifier[2].in_features, 7)
+
+        state_dict = torch.load(weights_path, map_location=self.device)
+        self.model.load_state_dict(state_dict)
+        self.model.eval()
+
+        self.transform = ConvNeXt_Base_Weights.IMAGENET1K_V1.transforms()
+
+    def make_prediction(self, image: Image):
+        image = self.transform(image).unsqueeze(0)
+        with torch.no_grad():
+            pred = F.softmax(self.model(image), dim=1)
+        return pred
+
+    

HAM_10000_CLASSIFICATION-master/models/efficient_net.py

@@ -0,0 +1,31 @@
+import torch
+from torchvision.models import efficientnet_b7
+from torchvision.transforms import InterpolationMode
+from torchvision import transforms
+from PIL import Image
+from torch.nn import functional as F
+
+class EfficientNetB7:
+    def __init__(self, weights_path: str):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        self.transform = transforms.Compose([
+            transforms.Resize((600), interpolation=InterpolationMode.BICUBIC),
+            transforms.CenterCrop(600),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                std=[0.229, 0.224, 0.225])
+        ])
+
+        self.model = efficientnet_b7()
+        self.model.classifier[1] = torch.nn.Linear(2560, 7)
+
+        state_dict = torch.load(weights_path, map_location=torch.device(self.device))
+        self.model.load_state_dict(state_dict)
+        self.model.eval()
+    
+    def make_prediction(self, image: Image):
+        image = self.transform(image).unsqueeze(0)
+        with torch.no_grad():
+            pred = F.softmax(self.model(image), dim=1)
+        return pred

HAM_10000_CLASSIFICATION-master/models/mobile_net.py

@@ -0,0 +1,39 @@
+import torch
+from torchvision.models import mobilenet_v3_large, MobileNet_V3_Large_Weights
+from torchvision.transforms import InterpolationMode
+from torchvision import transforms
+from PIL import Image
+from torch.nn import functional as F
+
+class MobileNetV3Large:
+    def __init__(self, weights_path: str):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        self.transform = transforms.Compose([
+            transforms.Resize((232), interpolation=InterpolationMode.BILINEAR),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                std=[0.229, 0.224, 0.225])
+        ])
+
+        self.model = mobilenet_v3_large()
+        self.model.classifier = torch.nn.Sequential(
+            torch.nn.Linear(960, 1280),
+            torch.nn.ReLU(),
+            torch.nn.Dropout(0.2),
+            torch.nn.Linear(1280, 1000),
+            torch.nn.ReLU(),
+            torch.nn.Dropout(0.2),
+            torch.nn.Linear(1000, 7)
+        )
+
+        state_dict = torch.load(weights_path, map_location=torch.device(self.device))
+        self.model.load_state_dict(state_dict)
+        self.model.eval()
+    
+    def make_prediction(self, image: Image):
+        image = self.transform(image).unsqueeze(0)
+        with torch.no_grad():
+            pred = F.softmax(self.model(image), dim=1)
+        return pred

HAM_10000_CLASSIFICATION-master/models/shufflenet.py

@@ -0,0 +1,25 @@
+import torch
+from torchvision.models import shufflenet_v2_x1_5, ShuffleNet_V2_X1_5_Weights
+from PIL import Image
+import torch.nn.functional as F
+
+class ShuffleNet_V2_X1_5:
+    def __init__(self, weights_path: str):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        self.model = shufflenet_v2_x1_5()
+        self.model.fc = torch.nn.Linear(self.model.fc.in_features, 7)
+
+        state_dict = torch.load(weights_path, map_location=self.device)
+        self.model.load_state_dict(state_dict)
+        self.model.eval()
+
+        self.transform = ShuffleNet_V2_X1_5_Weights.IMAGENET1K_V1.transforms()
+
+    def make_prediction(self, image: Image):
+        image = self.transform(image).unsqueeze(0)
+        with torch.no_grad():
+            pred = F.softmax(self.model(image), dim=1)
+        return pred
+
+    

HAM_10000_CLASSIFICATION-master/models/vit.py

@@ -0,0 +1,31 @@
+import torch
+from torchvision.models import vit_b_16, ViT_B_16_Weights
+from PIL import Image
+import torch.nn.functional as F
+from torchvision import transforms
+
+class VIT_B_16:
+    def __init__(self, weights_path: str):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        self.model = vit_b_16()
+        self.model.heads[0] = torch.nn.Linear(768, 7)
+
+        state_dict = torch.load(weights_path, map_location=self.device)
+        self.model.load_state_dict(state_dict)
+        self.model.eval()
+        
+        self.transform = transforms.Compose([
+            transforms.Resize(256),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ])
+
+    def make_prediction(self, image: Image):
+        image = self.transform(image).unsqueeze(0)
+        with torch.no_grad():
+            pred = F.softmax(self.model(image), dim=1)
+        return pred
+
+    

HAM_10000_CLASSIFICATION-master/requirements.txt

@@ -0,0 +1,83 @@
+altair==5.5.0
+annotated-types==0.7.0
+anyio==4.9.0
+attrs==25.3.0
+beautifulsoup4==4.13.3
+blinker==1.9.0
+cachetools==5.5.2
+certifi==2025.1.31
+charset-normalizer==3.4.1
+click==8.1.8
+colorama==0.4.6
+contourpy==1.3.1
+cycler==0.12.1
+fastapi==0.115.12
+filelock==3.18.0
+fonttools==4.56.0
+fsspec==2025.3.2
+gdown==5.2.0
+gitdb==4.0.12
+GitPython==3.1.44
+h11==0.14.0
+idna==3.10
+Jinja2==3.1.6
+jsonschema==4.23.0
+jsonschema-specifications==2024.10.1
+kiwisolver==1.4.8
+MarkupSafe==3.0.2
+matplotlib==3.10.1
+mpmath==1.3.0
+narwhals==1.32.0
+networkx==3.4.2
+numpy==2.2.4
+nvidia-cublas-cu12==12.4.5.8
+nvidia-cuda-cupti-cu12==12.4.127
+nvidia-cuda-nvrtc-cu12==12.4.127
+nvidia-cuda-runtime-cu12==12.4.127
+nvidia-cudnn-cu12==9.1.0.70
+nvidia-cufft-cu12==11.2.1.3
+nvidia-curand-cu12==10.3.5.147
+nvidia-cusolver-cu12==11.6.1.9
+nvidia-cusparse-cu12==12.3.1.170
+nvidia-cusparselt-cu12==0.6.2
+nvidia-nccl-cu12==2.21.5
+nvidia-nvjitlink-cu12==12.4.127
+nvidia-nvtx-cu12==12.4.127
+opencv-python-headless==4.11.0.86
+packaging==24.2
+pandas==2.2.3
+pillow==11.1.0
+protobuf==5.29.4
+pyarrow==19.0.1
+pydantic==2.11.1
+pydantic_core==2.33.0
+pydeck==0.9.1
+pyparsing==3.2.3
+PySocks==1.7.1
+python-dateutil==2.9.0.post0
+python-multipart==0.0.20
+pytz==2025.2
+referencing==0.36.2
+requests==2.32.3
+rpds-py==0.24.0
+setuptools==78.1.0
+six==1.17.0
+smmap==5.0.2
+sniffio==1.3.1
+soupsieve==2.6
+starlette==0.46.1
+streamlit==1.44.0
+sympy==1.13.1
+tenacity==9.0.0
+toml==0.10.2
+torch==2.6.0
+torchvision==0.21.0
+tornado==6.4.2
+tqdm==4.67.1
+triton==3.2.0
+typing-inspection==0.4.0
+typing_extensions==4.13.0
+tzdata==2025.2
+urllib3==2.3.0
+uvicorn==0.34.0
+watchdog==6.0.0

HAM_10000_CLASSIFICATION-master/utils/trainer.py

@@ -0,0 +1,216 @@
+def train_model(model, train_loader, val_loader, test_loader, optimizer, criterion, epochs, output_folder, device="cpu"):
+    """
+    Train a neural network model with specified training, validation, and testing datasets.
+    Additionally, plots accuracy and loss per epoch using matplotlib and saves them as images.
+
+    This function performs a complete training loop, including:
+    - Creating DataLoaders for training, validation, and testing datasets
+    - Moving the model to the specified device (CPU/GPU)
+    - Training the model for a specified number of epochs
+    - Tracking and logging training, validation, and testing metrics
+    - Saving the best (based on validation performance) and last model weights
+    - Plotting and saving accuracy and loss graphs per epoch
+
+    Parameters:
+    -----------
+    model : torch.nn.Module
+        The neural network model to be trained
+    train_loader : torch.utils.data.DataLoader
+        Dataset used for training the model
+    val_loader : torch.utils.data.DataLoader
+        Dataset used for validating the model during training
+    test_loader : torch.utils.data.DataLoader
+        Dataset used for evaluating the model's performance after training
+    optimizer : torch.optim.Optimizer
+        Optimization algorithm for updating model weights
+    criterion : torch.nn.Module
+        Loss function used to compute the model's performance
+    epochs : int
+        Number of complete passes through the entire training dataset
+    output_folder : str
+        Folder path where the model weights and plots will be saved
+    device : str, optional
+        Computing device to use for training (default is "cpu")
+        Can be "cpu" or "cuda" for GPU training
+
+    Returns:
+    --------
+    None
+
+    Side Effects:
+    -------------
+    - Prints training, validation, and testing metrics for each epoch
+    - Saves the best performing model (based on validation accuracy) to "weights/best_model.pth"
+    - Saves the final model to "weights/last_model.pth"
+    - Saves the loss plot as "loss_plot.png" and accuracy plot as "accuracy_plot.png" in the output folder
+
+    Example:
+    --------
+    >>> model = MyModel()
+    >>> optimizer = torch.optim.Adam(model.parameters())
+    >>> criterion = nn.CrossEntropyLoss()
+    >>> train_model(model, train_loader, val_loader, test_loader, optimizer, criterion, epochs=10, batch_size=32, output_folder="weights")
+    """
+    import os
+    import torch
+    from torch.utils.data import DataLoader
+    from tqdm import tqdm
+    import matplotlib.pyplot as plt
+
+    # Ensure weights folder exists
+    os.makedirs(output_folder, exist_ok=True)
+
+    print(f"Device Found: {device}, Starting Training 🚀")
+    
+    # Move model to the specified device
+    model = model.to(device)
+
+    best_val_accuracy = 0.0  # Initialize best validation accuracy tracker
+
+    # Lists to store metrics per epoch for plotting
+    train_losses, val_losses, test_losses = [], [], []
+    train_accuracies, val_accuracies, test_accuracies = [], [], []
+
+    for epoch in range(epochs):
+        # ----------------------
+        # Training Phase
+        # ----------------------
+        model.train()  # Set model to training mode
+        running_loss = 0.0
+        correct = 0
+        total = 0
+
+        train_progress = tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs} (Training)", leave=False)
+        for images, labels in train_progress:
+            # Move tensors to the specified device
+            images, labels = images.to(device), labels.to(device)
+
+            optimizer.zero_grad()  # Reset gradients
+            outputs = model(images)  # Forward pass
+            loss = criterion(outputs, labels)  # Compute loss
+            loss.backward()  # Backpropagation
+            optimizer.step()  # Update weights
+
+            running_loss += loss.item()
+            _, predicted = torch.max(outputs, 1)
+            total += labels.size(0)
+            correct += (predicted == labels).sum().item()
+
+            train_progress.set_postfix({
+                'Loss': f'{loss.item():.4f}', 
+                'Accuracy': f'{100 * correct / total:.2f}%'
+            })
+
+        train_loss = running_loss / len(train_loader)
+        train_accuracy = 100 * correct / total
+
+        # ----------------------
+        # Validation Phase
+        # ----------------------
+        model.eval()  # Set model to evaluation mode
+        val_loss = 0.0
+        correct_val = 0
+        total_val = 0
+
+        val_progress = tqdm(val_loader, desc=f"Epoch {epoch+1}/{epochs} (Validation)", leave=False)
+        with torch.no_grad():
+            for images, labels in val_progress:
+                images, labels = images.to(device), labels.to(device)
+                outputs = model(images)
+                loss = criterion(outputs, labels)
+                val_loss += loss.item()
+                _, predicted = torch.max(outputs, 1)
+                total_val += labels.size(0)
+                correct_val += (predicted == labels).sum().item()
+
+                val_progress.set_postfix({
+                    'Loss': f'{loss.item():.4f}', 
+                    'Accuracy': f'{100 * correct_val / total_val:.2f}%'
+                })
+
+        val_loss /= len(val_loader)
+        val_accuracy = 100 * correct_val / total_val
+
+        # ----------------------
+        # Testing Phase
+        # ----------------------
+        test_loss = 0.0
+        correct_test = 0
+        total_test = 0
+
+        test_progress = tqdm(test_loader, desc=f"Epoch {epoch+1}/{epochs} (Testing)", leave=False)
+        with torch.no_grad():
+            for images, labels in test_progress:
+                images, labels = images.to(device), labels.to(device)
+                outputs = model(images)
+                loss = criterion(outputs, labels)
+                test_loss += loss.item()
+                _, predicted = torch.max(outputs, 1)
+                total_test += labels.size(0)
+                correct_test += (predicted == labels).sum().item()
+
+                test_progress.set_postfix({
+                    'Loss': f'{loss.item():.4f}', 
+                    'Accuracy': f'{100 * correct_test / total_test:.2f}%'
+                })
+
+        test_loss /= len(test_loader)
+        test_accuracy = 100 * correct_test / total_test
+
+        # Store metrics for plotting
+        train_losses.append(train_loss)
+        val_losses.append(val_loss)
+        test_losses.append(test_loss)
+        train_accuracies.append(train_accuracy)
+        val_accuracies.append(val_accuracy)
+        test_accuracies.append(test_accuracy)
+
+        # Log the metrics for this epoch
+        print(
+            f"Epoch [{epoch+1}/{epochs}]: "
+            f"Train Loss: {train_loss:.4f}, Train Accuracy: {train_accuracy:.2f}% | "
+            f"Validation Loss: {val_loss:.4f}, Validation Accuracy: {val_accuracy:.2f}% | "
+            f"Test Loss: {test_loss:.4f}, Test Accuracy: {test_accuracy:.2f}%"
+        )
+
+        # Save the best model based on validation accuracy
+        if val_accuracy > best_val_accuracy:
+            best_val_accuracy = val_accuracy
+            torch.save(model.state_dict(), os.path.join(output_folder, "best_model.pth"))
+
+    # Save the last model
+    torch.save(model.state_dict(), os.path.join(output_folder, "last_model.pth"))
+    print("Training completed. Best validation accuracy: {:.2f}%".format(best_val_accuracy))
+    
+    # ----------------------
+    # Plotting Metrics with Matplotlib
+    # ----------------------
+    epochs_range = range(1, epochs + 1)
+    
+    # Plot Losses
+    plt.figure()
+    plt.plot(epochs_range, train_losses, label='Train Loss')
+    plt.plot(epochs_range, val_losses, label='Validation Loss')
+    plt.plot(epochs_range, test_losses, label='Test Loss')
+    plt.xlabel('Epoch')
+    plt.ylabel('Loss')
+    plt.title('Loss per Epoch')
+    plt.legend()
+    loss_plot_path = os.path.join(output_folder, 'loss_plot.png')
+    plt.savefig(loss_plot_path)
+    plt.close()
+    print(f"Loss plot saved to {loss_plot_path}")
+
+    # Plot Accuracies
+    plt.figure()
+    plt.plot(epochs_range, train_accuracies, label='Train Accuracy')
+    plt.plot(epochs_range, val_accuracies, label='Validation Accuracy')
+    plt.plot(epochs_range, test_accuracies, label='Test Accuracy')
+    plt.xlabel('Epoch')
+    plt.ylabel('Accuracy (%)')
+    plt.title('Accuracy per Epoch')
+    plt.legend()
+    acc_plot_path = os.path.join(output_folder, 'accuracy_plot.png')
+    plt.savefig(acc_plot_path)
+    plt.close()
+    print(f"Accuracy plot saved to {acc_plot_path}")