Upload 9 files

.gitignore

@@ -0,0 +1,7 @@
+.vscode
+pokemonclassification 
+*.zip
+trained_models/pokemon_vgg.pth
+*.cpython-310.pyc
+__pycache__
+utils/__pycache__/data.cpython-310.pyc

README.md

@@ -1,3 +1,135 @@
----
-license: mit
----
+# PokemonClassification
+
+This repository explores the training of different models for a vision classification task, with a special focus made on reproducibility, and an attempt to a local interpretability of the decision made by a resnet model using LIME
+
+## Table of Contents
+
+- [PokemonClassification](#pokemonclassification)
+  - [Table of Contents](#table-of-contents)
+  - [Installation](#installation)
+  - [Dataset](#dataset)
+  - [Training](#training)
+  - [Inference](#inference)
+  - [Generating Data Samples](#generating-data-samples)
+  - [Interpretability](#interpretability)
+  - [Contributing](#contributing)
+
+## Installation
+
+1. Clone the repository:
+
+    ```sh
+    git clone https://github.com/yourusername/PokemonClassification.git
+    cd PokemonClassification
+    ```
+
+2. Create a conda environment and activate it:
+
+    ```sh
+    conda env create -f environment.yaml
+    conda activate pokemonclassification
+    ```
+
+## Dataset
+
+To get the data, use the appropriate script based on your operating system:
+
+- On Linux-based systems:
+  
+    ```shell
+    ./get_data.sh
+    ```
+
+- On Windows:
+  
+    ```shell
+    ./get_data.ps1
+    ```
+
+## Training
+
+To train a model, use the `train.py` script. Here are the parameters you can specify:
+
+```python
+def parser_args():
+    parser = argparse.ArgumentParser(description="Pokemon Classification")
+    parser.add_argument("--data_dir", type=str, default="./pokemonclassification/PokemonData", help="Path to the data directory")
+    parser.add_argument("--indices_file", type=str, default="indices_60_32.pkl", help="Path to the indices file")
+    parser.add_argument("--epochs", type=int, default=20, help="Number of epochs")
+    parser.add_argument("--lr", type=float, default=0.001, help="Learning rate")
+    parser.add_argument("--train_batch_size", type=int, default=128, help="train Batch size")
+    parser.add_argument("--test_batch_size", type=int, default=512, help="test Batch size")
+    parser.add_argument("--model", type=str, choices=["resnet", "alexnet", "vgg", "squeezenet", "densenet"], default="resnet", help="Model to be used")
+    parser.add_argument("--feature_extract", type=bool, default=True, help="whether to freeze the backbone or not")
+    parser.add_argument("--use_pretrained", type=bool, default=True, help="whether to use pretrained model or not")
+    parser.add_argument("--experiment_id", type=int, default=0, help="Experiment ID to log the results")
+    return parser.parse_args()
+```
+
+Example:
+
+```shell
+python train.py--model resnet --data_dir data/PokemonData --epochs 10 --train_batch_size 32 --test_batch_size 32
+```
+
+## Inference
+
+To perform inference on a single image, use the `inference.py` script. Here are the parameters you can specify:
+
+```python
+def main():
+    parser = argparse.ArgumentParser(description="Image Inference")
+    parser.add_argument("--model_name", type=str, help="Model name (resnet, alexnet, vgg, squeezenet, densenet)", default="resnet")
+    parser.add_argument("--model_weights", type=str, help="Path to the model weights", default="./trained_models/pokemon_resnet.pth")
+    parser.add_argument("--image_path", type=str, help="Path to the image", default="./pokemonclassification/PokemonData/Chansey/57ccf27cba024fac9531baa9f619ec62.jpg")
+    parser.add_argument("--num_classes", type=int, help="Number of classes", default=150)
+    parser.add_argument("--lime_interpretability", action="store_true", help="Whether to run interpretability or not")
+    parser.add_argument("--classify", action="store_true", help="Whether to classify the image when saving the lime filter")
+    args = parser.parse_args()
+
+    if args.lime_interpretability:
+        assert args.model_name == "resnet", "Interpretability is only supported for ResNet model for now"
+```
+
+Example:
+
+```shell
+python inference.py --model_name resnet --model_weights path_to_your_model_weights.pth --image_path path_to_your_image.jpg --num_classes 10
+```
+
+## Generating Data Samples
+
+To generate data samples, use the `get_samples.py` script. Here are the parameters you can specify:
+
+```python
+def main():
+    parser = argparse.ArgumentParser(description="Generate Data Samples")
+    parser.add_argument("--model_name", type=str, help="Model name (resnet, alexnet, vgg, squeezenet, densenet)", default="resnet")
+    parser.add_argument("--model_weights", type=str, help="Path to the model weights", default="./trained_models/pokemon_resnet.pth")
+    parser.add_argument("--image_path", type=str, help="Path to the image", default="./pokemonclassification/PokemonData/")
+    parser.add_argument("--num_classes", type=int, help="Number of classes", default=150)
+    parser.add_argument("--label", type=str, help="Label to filter the images", default='Dragonair')
+    parser.add_argument("--num_correct", type=int, help="Number of correctly classified images", default=5)
+    parser.add_argument("--num_incorrect", type=int, help="Number of incorrectly classified images", default=5)
+    args = parser.parse_args()
+```
+
+Example:
+
+```shell
+python get_samples.py --model_name resnet --model_weights path_to_your_model_weights.pth --image_path path_to_your_image_directory --num_classes 10 --label Pikachu --num_correct 5 --num_incorrect 5
+```
+
+## Interpretability
+
+To interpret the model's predictions using LIME, use the `inference.py` script with the `--lime_interpretability` flag.
+
+Example:
+
+```shell
+python inference.py --model_name resnet --model_weights path_to_your_model_weights.pth --image_path path_to_your_image.jpg --num_classes 10 --lime_interpretability
+```
+
+## Contributing
+
+Contributions are welcome! Please open an issue or submit a pull request for any improvements or bug fixes.

environment.yaml

@@ -0,0 +1,313 @@
+name: cloudspace
+channels:
+  - conda-forge
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1=main
+  - _openmp_mutex=5.1=1_gnu
+  - aiohappyeyeballs=2.4.3=py310h06a4308_0
+  - alembic=1.13.3=py310h06a4308_0
+  - aniso8601=9.0.1=pyhd3eb1b0_0
+  - arrow-cpp=16.1.0=hc1eb8f0_0
+  - attrs=24.2.0=py310h06a4308_0
+  - bcrypt=3.2.0=py310h5eee18b_1
+  - blas=1.0=openblas
+  - blinker=1.6.2=py310h06a4308_0
+  - boost-cpp=1.82.0=hdb19cb5_2
+  - bottleneck=1.4.2=py310ha9d4c09_0
+  - brotli=1.0.9=h5eee18b_8
+  - brotli-bin=1.0.9=h5eee18b_8
+  - brotli-python=1.0.9=py310h6a678d5_8
+  - bzip2=1.0.8=h5eee18b_6
+  - c-ares=1.19.1=h5eee18b_0
+  - ca-certificates=2024.11.26=h06a4308_0
+  - certifi=2024.8.30=py310h06a4308_0
+  - cffi=1.17.1=py310h1fdaa30_0
+  - click=8.1.7=py310h06a4308_0
+  - cloudpickle=3.0.0=py310h06a4308_0
+  - contourpy=1.3.1=py310hdb19cb5_0
+  - cryptography=43.0.3=py310h7825ff9_1
+  - databricks-sdk=0.33.0=py310h06a4308_0
+  - deprecated=1.2.13=py310h06a4308_0
+  - docker-py=7.1.0=py310h06a4308_0
+  - entrypoints=0.4=py310h06a4308_0
+  - flask=3.0.3=py310h06a4308_0
+  - freetype=2.12.1=h4a9f257_0
+  - frozenlist=1.5.0=py310h5eee18b_0
+  - gflags=2.2.2=h6a678d5_1
+  - gitdb=4.0.7=pyhd3eb1b0_0
+  - gitpython=3.1.43=py310h06a4308_0
+  - glog=0.5.0=h6a678d5_1
+  - graphene=3.3=py310h06a4308_0
+  - graphql-core=3.2.3=py310h06a4308_1
+  - graphql-relay=3.2.0=py310h06a4308_0
+  - greenlet=3.0.1=py310h6a678d5_0
+  - gunicorn=22.0.0=py310h06a4308_0
+  - icu=73.1=h6a678d5_0
+  - importlib-metadata=8.5.0=py310h06a4308_0
+  - itsdangerous=2.2.0=py310h06a4308_0
+  - jinja2=3.1.4=py310h06a4308_1
+  - joblib=1.4.2=py310h06a4308_0
+  - jpeg=9e=h5eee18b_3
+  - krb5=1.20.1=h143b758_1
+  - lcms2=2.12=h3be6417_0
+  - ld_impl_linux-64=2.40=h12ee557_0
+  - lerc=3.0=h295c915_0
+  - libabseil=20240116.2=cxx17_h6a678d5_0
+  - libboost=1.82.0=h109eef0_2
+  - libbrotlicommon=1.0.9=h5eee18b_8
+  - libbrotlidec=1.0.9=h5eee18b_8
+  - libbrotlienc=1.0.9=h5eee18b_8
+  - libcurl=8.9.1=h251f7ec_0
+  - libdeflate=1.17=h5eee18b_1
+  - libedit=3.1.20230828=h5eee18b_0
+  - libev=4.33=h7f8727e_1
+  - libevent=2.1.12=hdbd6064_1
+  - libffi=3.4.4=h6a678d5_1
+  - libgcc-ng=11.2.0=h1234567_1
+  - libgfortran-ng=11.2.0=h00389a5_1
+  - libgfortran5=11.2.0=h1234567_1
+  - libgomp=11.2.0=h1234567_1
+  - libgrpc=1.62.2=h2d74bed_0
+  - libnghttp2=1.57.0=h2d74bed_0
+  - libopenblas=0.3.21=h043d6bf_0
+  - libpng=1.6.39=h5eee18b_0
+  - libprotobuf=4.25.3=he621ea3_0
+  - libsodium=1.0.18=h7b6447c_0
+  - libssh2=1.11.1=h251f7ec_0
+  - libstdcxx-ng=11.2.0=h1234567_1
+  - libthrift=0.15.0=h1795dd8_2
+  - libtiff=4.5.1=h6a678d5_0
+  - libuuid=1.41.5=h5eee18b_0
+  - libwebp-base=1.3.2=h5eee18b_1
+  - lz4-c=1.9.4=h6a678d5_1
+  - mako=1.2.3=py310h06a4308_0
+  - matplotlib-base=3.9.2=py310hbfdbfaf_1
+  - mlflow=2.18.0=hff52083_0
+  - mlflow-skinny=2.18.0=py310hff52083_0
+  - mlflow-ui=2.18.0=py310hff52083_0
+  - multidict=6.1.0=py310h5eee18b_0
+  - ncurses=6.4=h6a678d5_0
+  - numexpr=2.10.1=py310hd28fd6d_0
+  - numpy=1.26.4=py310heeff2f4_0
+  - numpy-base=1.26.4=py310h8a23956_0
+  - openjpeg=2.5.2=he7f1fd0_0
+  - openssl=3.0.15=h5eee18b_0
+  - opentelemetry-api=1.16.0=pyhd8ed1ab_0
+  - opentelemetry-sdk=1.16.0=pyhd8ed1ab_0
+  - opentelemetry-semantic-conventions=0.37b0=pyhd8ed1ab_0
+  - orc=2.0.1=h2d29ad5_0
+  - paramiko=3.5.0=py310h06a4308_0
+  - pillow=11.0.0=py310hfdbf927_0
+  - prometheus_client=0.21.0=py310h06a4308_0
+  - prometheus_flask_exporter=0.22.4=py310h06a4308_0
+  - propcache=0.2.0=py310h5eee18b_0
+  - pyarrow=16.1.0=py310h1128e8f_0
+  - pynacl=1.5.0=py310h5eee18b_0
+  - pyopenssl=24.2.1=py310h06a4308_0
+  - pyparsing=3.2.0=py310h06a4308_0
+  - pysocks=1.7.1=py310h06a4308_0
+  - python=3.10.15=he870216_1
+  - python-dateutil=2.9.0post0=py310h06a4308_2
+  - python-tzdata=2023.3=pyhd3eb1b0_0
+  - python_abi=3.10=2_cp310
+  - pyyaml=6.0.2=py310h5eee18b_0
+  - querystring_parser=1.2.4=py310h06a4308_0
+  - re2=2022.04.01=h295c915_0
+  - readline=8.2=h5eee18b_0
+  - requests=2.32.3=py310h06a4308_1
+  - s2n=1.3.27=hdbd6064_0
+  - setuptools=75.1.0=py310h06a4308_0
+  - six=1.16.0=pyhd3eb1b0_1
+  - smmap=4.0.0=pyhd3eb1b0_0
+  - snappy=1.2.1=h6a678d5_0
+  - sqlalchemy=2.0.34=py310h00e1ef3_0
+  - sqlite=3.45.3=h5eee18b_0
+  - sqlparse=0.4.4=py310h06a4308_0
+  - threadpoolctl=3.5.0=py310h2f386ee_0
+  - tk=8.6.14=h39e8969_0
+  - typing_extensions=4.11.0=py310h06a4308_0
+  - unicodedata2=15.1.0=py310h5eee18b_0
+  - urllib3=2.2.3=py310h06a4308_0
+  - utf8proc=2.6.1=h5eee18b_1
+  - websocket-client=1.8.0=py310h06a4308_0
+  - wheel=0.44.0=py310h06a4308_0
+  - wrapt=1.14.1=py310h5eee18b_0
+  - xz=5.4.6=h5eee18b_1
+  - yaml=0.2.5=h7b6447c_0
+  - yarl=1.18.0=py310h5eee18b_0
+  - zipp=3.21.0=py310h06a4308_0
+  - zlib=1.2.13=h5eee18b_1
+  - zstd=1.5.6=hc292b87_0
+  - pip:
+      - absl-py==2.1.0
+      - aiohttp==3.11.7
+      - aiosignal==1.3.1
+      - annotated-types==0.7.0
+      - anyio==4.6.2.post1
+      - argon2-cffi==23.1.0
+      - argon2-cffi-bindings==21.2.0
+      - arrow==1.3.0
+      - asttokens==2.4.1
+      - async-lru==2.0.4
+      - async-timeout==5.0.1
+      - babel==2.16.0
+      - backoff==2.2.1
+      - beautifulsoup4==4.12.3
+      - bleach==6.2.0
+      - boto3==1.35.70
+      - botocore==1.35.70
+      - cachetools==5.5.0
+      - charset-normalizer==3.4.0
+      - comm==0.2.2
+      - cycler==0.12.1
+      - debugpy==1.8.9
+      - decorator==5.1.1
+      - defusedxml==0.7.1
+      - exceptiongroup==1.2.2
+      - executing==2.1.0
+      - fastapi==0.115.5
+      - fastjsonschema==2.20.0
+      - filelock==3.16.1
+      - fire==0.7.0
+      - fonttools==4.55.0
+      - fqdn==1.5.1
+      - fsspec==2024.10.0
+      - git-filter-repo==2.47.0
+      - google-auth==2.36.0
+      - google-auth-oauthlib==1.2.1
+      - grpcio==1.68.0
+      - h11==0.14.0
+      - httpcore==1.0.7
+      - httptools==0.6.4
+      - httpx==0.27.2
+      - idna==3.10
+      - imageio==2.36.1
+      - ipykernel==6.26.0
+      - ipython==8.17.2
+      - ipywidgets==8.1.1
+      - isoduration==20.11.0
+      - jedi==0.19.2
+      - jmespath==1.0.1
+      - json5==0.10.0
+      - jsonpointer==3.0.0
+      - jsonschema==4.23.0
+      - jsonschema-specifications==2024.10.1
+      - jupyter-client==8.6.3
+      - jupyter-core==5.7.2
+      - jupyter-events==0.10.0
+      - jupyter-lsp==2.2.5
+      - jupyter-server==2.14.2
+      - jupyter-server-terminals==0.5.3
+      - jupyterlab==4.2.0
+      - jupyterlab-pygments==0.3.0
+      - jupyterlab-server==2.27.3
+      - jupyterlab-widgets==3.0.13
+      - kiwisolver==1.4.7
+      - lazy-loader==0.4
+      - lightning==2.4.0
+      - lightning-cloud==0.5.70
+      - lightning-sdk==0.1.30
+      - lightning-utilities==0.11.9
+      - lime==0.2.0.1
+      - litdata==0.2.32
+      - litserve==0.2.5
+      - markdown==3.7
+      - markdown-it-py==3.0.0
+      - markupsafe==3.0.2
+      - matplotlib==3.8.2
+      - matplotlib-inline==0.1.7
+      - mdurl==0.1.2
+      - mistune==3.0.2
+      - mpmath==1.3.0
+      - nbclient==0.10.0
+      - nbconvert==7.16.4
+      - nbformat==5.10.4
+      - nest-asyncio==1.6.0
+      - networkx==3.4.2
+      - notebook-shim==0.2.4
+      - nvidia-cublas-cu12==12.1.3.1
+      - nvidia-cuda-cupti-cu12==12.1.105
+      - nvidia-cuda-nvrtc-cu12==12.1.105
+      - nvidia-cuda-runtime-cu12==12.1.105
+      - nvidia-cudnn-cu12==8.9.2.26
+      - nvidia-cufft-cu12==11.0.2.54
+      - nvidia-curand-cu12==10.3.2.106
+      - nvidia-cusolver-cu12==11.4.5.107
+      - nvidia-cusparse-cu12==12.1.0.106
+      - nvidia-nccl-cu12==2.19.3
+      - nvidia-nvjitlink-cu12==12.6.85
+      - nvidia-nvtx-cu12==12.1.105
+      - oauthlib==3.2.2
+      - overrides==7.7.0
+      - packaging==24.2
+      - pandas==2.1.4
+      - pandocfilters==1.5.1
+      - parso==0.8.4
+      - pexpect==4.9.0
+      - pip==24.3.1
+      - platformdirs==4.3.6
+      - prompt-toolkit==3.0.48
+      - protobuf==4.23.4
+      - psutil==6.1.0
+      - ptyprocess==0.7.0
+      - pure-eval==0.2.3
+      - pyasn1==0.6.1
+      - pyasn1-modules==0.4.1
+      - pycparser==2.22
+      - pydantic==2.10.2
+      - pydantic-core==2.27.1
+      - pygments==2.18.0
+      - pyjwt==2.10.0
+      - python-dotenv==1.0.1
+      - python-json-logger==2.0.7
+      - python-multipart==0.0.17
+      - pytorch-lightning==2.4.0
+      - pytz==2024.2
+      - pyzmq==26.2.0
+      - referencing==0.35.1
+      - requests-oauthlib==2.0.0
+      - rfc3339-validator==0.1.4
+      - rfc3986-validator==0.1.1
+      - rich==13.9.4
+      - rpds-py==0.21.0
+      - rsa==4.9
+      - s3transfer==0.10.4
+      - scikit-image==0.24.0
+      - scikit-learn==1.3.2
+      - scipy==1.11.4
+      - send2trash==1.8.3
+      - simple-term-menu==1.6.5
+      - sniffio==1.3.1
+      - soupsieve==2.6
+      - stack-data==0.6.3
+      - starlette==0.41.3
+      - sympy==1.13.3
+      - tensorboard==2.15.1
+      - tensorboard-data-server==0.7.2
+      - termcolor==2.5.0
+      - terminado==0.18.1
+      - tifffile==2024.9.20
+      - tinycss2==1.4.0
+      - tomli==2.1.0
+      - torch==2.2.1+cu121
+      - torchmetrics==1.3.1
+      - torchvision==0.17.1+cu121
+      - tornado==6.4.2
+      - tqdm==4.67.1
+      - traitlets==5.14.3
+      - triton==2.2.0
+      - types-python-dateutil==2.9.0.20241003
+      - typing-extensions==4.12.2
+      - tzdata==2024.2
+      - uri-template==1.3.0
+      - uvicorn==0.32.1
+      - uvloop==0.21.0
+      - watchfiles==1.0.0
+      - wcwidth==0.2.13
+      - webcolors==24.11.1
+      - webencodings==0.5.1
+      - websockets==14.1
+      - werkzeug==3.1.3
+      - widgetsnbextension==4.0.13
+prefix: /home/zeus/miniconda3/envs/cloudspace

get_data.ps1

@@ -0,0 +1,11 @@
+# Download the file using Invoke-WebRequest
+$destination = "./pokemonclassification.zip"
+$url = "https://www.kaggle.com/api/v1/datasets/download/lantian773030/pokemonclassification"
+Invoke-WebRequest -Uri $url -OutFile $destination
+
+# Extract the zip file to the specified folder
+$extractPath = "./pokemonclassification"
+Expand-Archive -Path $destination -DestinationPath $extractPath
+
+# Remove the downloaded zip file (if not needed anymore)
+Remove-Item $destination

get_data.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+curl -L -o ./pokemonclassification.zip https://www.kaggle.com/api/v1/datasets/download/lantian773030/pokemonclassification
+
+unzip -d ./pokemonclassification pokemonclassification.zip
+
+# Remove the downloaded zip file (if you don't need it anymore)
+rm ./pokemonclassification.zip

get_samples.py

@@ -0,0 +1,55 @@
+from utils.inference_utils import find_images_from_path
+import torch
+import argparse
+from utils.train_utils import initialize_model
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Image Inference")
+    parser.add_argument(
+        "--model_name",
+        type=str,
+        help="Model name (resnet, alexnet, vgg, squeezenet, densenet)",
+        default="resnet",
+    )
+    parser.add_argument(
+        "--model_weights",
+        type=str,
+        help="Path to the model weights",
+        default="./trained_models/pokemon_resnet.pth",
+    )
+    parser.add_argument(
+        "--image_path",
+        type=str,
+        help="Path to the image",
+        default="./pokemonclassification/PokemonData/",
+    )
+    parser.add_argument(
+        "--num_classes", type=int,  help="Number of classes", default=150
+    )
+    parser.add_argument(
+        "--label", type=str, help="Label to filter the images", default='Dragonair' # Krabby, Clefairy
+        )
+    parser.add_argument(
+        "--num_correct", type=int, help="Number of correctly classified images", default=5
+        )
+    parser.add_argument(
+        "--num_incorrect", type=int, help="Number of incorrectly classified images", default=5
+        )
+    
+    args = parser.parse_args()
+
+    assert (args.model_name == "resnet"), "Only the ResNet is supported model for now"
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Initialize the model
+    model = initialize_model(args.model_name, args.num_classes)
+    model = model.to(device)
+
+    # Load the model weights
+    model.load_state_dict(torch.load(args.model_weights, map_location=device))
+    find_images_from_path(args.image_path, model, device, args.num_correct, args.num_incorrect, args.label)
+
+if __name__ == "__main__":
+    main()

indices_60_32.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:32b56617770be9430d034b41ef9235bb938cd1db40a78bb15a7d579229c79511
+size 20240

inference.py

@@ -0,0 +1,70 @@
+import torch
+import argparse
+from utils.inference_utils import preprocess_image, predict
+from utils.train_utils import initialize_model
+from utils.interpretability import lime_interpret_image_inference
+from utils.data import CLASS_NAMES
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Image Inference")
+    parser.add_argument(
+        "--model_name",
+        type=str,
+        help="Model name (resnet, alexnet, vgg, squeezenet, densenet)",
+        default="resnet",
+    )
+    parser.add_argument(
+        "--model_weights",
+        type=str,
+        help="Path to the model weights",
+        default="./trained_models/pokemon_resnet.pth",
+    )
+    parser.add_argument(
+        "--image_path",
+        type=str,
+        help="Path to the image",
+        default="./pokemonclassification/PokemonData/Chansey/57ccf27cba024fac9531baa9f619ec62.jpg",
+    )
+    parser.add_argument(
+        "--num_classes", type=int,  help="Number of classes", default=150
+    )
+    parser.add_argument(
+        "--lime_interpretability",
+        action="store_true",
+        help="Whether to run interpretability or not",
+    )
+    parser.add_argument(
+        "--classify",
+        action="store_true",
+        help="Whether to classify the image when saving the lime filter")
+    
+    args = parser.parse_args()
+
+    if args.lime_interpretability:
+        assert (
+            args.model_name == "resnet"
+        ), "Interpretability is only supported for ResNet model for now"
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Initialize the model
+    model = initialize_model(args.model_name, args.num_classes)
+    model = model.to(device)
+
+    # Load the model weights
+    model.load_state_dict(torch.load(args.model_weights, map_location=device))
+
+    # Preprocess the image
+    image = preprocess_image(args.image_path, (224, 224)).to(device)
+
+    # Perform inference
+    preds = torch.max(predict(model, image), 1)[1]
+    print(f"Predicted class: {CLASS_NAMES[preds.item()]}")
+
+    if args.lime_interpretability:
+        lime_interpret_image_inference(args, model, image, device)
+
+
+if __name__ == "__main__":
+    main()

train.py

@@ -0,0 +1,174 @@
+import torch.nn as nn
+from torchvision import transforms
+from utils.data import PokemonDataModule
+from utils.train import initialize_model, train_and_evaluate
+import torch
+import torch.optim as optim
+import mlflow
+import argparse
+import random
+
+# The shape of the images that the models expects
+IMG_SHAPE = (224, 224)
+
+
+def parser_args():
+    parser = argparse.ArgumentParser(description="Pokemon Classification")
+    parser.add_argument(
+        "--data_dir",
+        type=str,
+        default="./pokemonclassification/PokemonData",
+        help="Path to the data directory",
+    )
+    parser.add_argument(
+        "--indices_file",
+        type=str,
+        default="indices_60_32.pkl",
+        help="Path to the indices file",
+    )
+    parser.add_argument("--epochs", type=int, default=20, help="Number of epochs")
+    parser.add_argument("--lr", type=float, default=0.001, help="Learning rate")
+    parser.add_argument(
+        "--train_batch_size", type=int, default=128, help="train Batch size"
+    )
+    parser.add_argument(
+        "--test_batch_size", type=int, default=512, help="test Batch size"
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        choices=["resnet", "alexnet", "vgg", "squeezenet", "densenet"],
+        default="resnet",
+        help="Model to be used",
+    )
+    parser.add_argument(
+        "--feature_extract",
+        type=bool,
+        default=True,
+        help="whether to freeze the backbone or not",
+    )
+    parser.add_argument(
+        "--use_pretrained",
+        type=bool,
+        default=True,
+        help="whether to use pretrained model or not",
+    )
+    parser.add_argument(
+        "--experiment_id",
+        type=int,
+        default=0,
+        help="Experiment ID to log the results",
+    )
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parser_args()
+
+    pokemon_dataset = PokemonDataModule(args.data_dir)
+    NUM_CLASSES = len(pokemon_dataset.class_names)
+
+    # Get class names
+    print(f"Number of classes: {NUM_CLASSES}")
+
+    # You can only the use precomputed means and vars if using the same indices file ('indices_60_32.pkl')
+    if "indices_60_32.pkl" in args.indices_file:
+        chanel_means = torch.tensor([0.6062, 0.5889, 0.5550])
+        chanel_vars = torch.tensor([0.3284, 0.3115, 0.3266])
+        stats = {"mean": chanel_means, "std": chanel_vars}
+        _ = pokemon_dataset.prepare_data(
+            indices_file=args.indices_file, get_stats=False
+        )
+    else:
+        stats = pokemon_dataset.prepare_data(
+            indices_file=args.indices_file, get_stats=True
+        )
+
+    print(f"Train dataset size: {len(pokemon_dataset.train_dataset)}")
+    print(f"Test dataset size: {len(pokemon_dataset.test_dataset)}")
+
+    # Transformations of data for testing
+    test_transform = transforms.Compose(
+        [
+            transforms.Resize(IMG_SHAPE),
+            transforms.ToTensor(),  # Convert PIL images to tensors
+            transforms.Normalize(**stats),  # Normalize images using mean and std
+        ]
+    )
+
+    # Data augmentations for training
+    train_transform = transforms.Compose(
+        [
+            transforms.Resize(IMG_SHAPE),
+            transforms.RandomRotation(10),
+            transforms.RandomHorizontalFlip(),
+            transforms.RandomCrop(IMG_SHAPE, padding=4),
+            transforms.RandomAffine(degrees=10, translate=(0.1, 0.1), scale=(0.9, 1.1)),
+            transforms.ColorJitter(
+                brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2
+            ),
+            transforms.RandomGrayscale(p=0.2),
+            transforms.ToTensor(),
+            transforms.Normalize(**stats),
+        ]
+    )
+
+    # get dataloaders
+    trainloader, testloader = pokemon_dataset.get_dataloaders(
+        train_transform=train_transform,
+        test_transform=test_transform,
+        train_batch_size=args.train_batch_size,
+        test_batch_size=args.test_batch_size,
+    )
+
+    pokemon_dataset.plot_examples(testloader, stats=stats)
+
+    pokemon_dataset.plot_examples(trainloader, stats=stats)
+
+    # Try with a finetuning a resnet for example
+    model = initialize_model(
+        args.model,
+        NUM_CLASSES,
+        feature_extract=args.feature_extract,
+        use_pretrained=args.use_pretrained,
+    )
+
+    # Print the model we just instantiated
+    print(model)
+
+    # Model, criterion, optimizer
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(model.parameters(), lr=args.lr)
+
+    # Device configuration
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    with mlflow.start_run(
+        experiment_id=args.experiment_id,
+        run_name=f"{args.model}_{'finetuning' if not args.feature_extract else 'feature_extracting'}"
+        f"_{'pretrained' if args.use_pretrained else 'not_pretrained'}"
+        f"_{args.indices_file}_{random.randint(0, 1000)}",
+    ) as run:
+        mlflow.log_param("epochs", args.epochs)
+        mlflow.log_param("lr", args.lr)
+        mlflow.log_param("train_batch_size", args.train_batch_size)
+        mlflow.log_param("test_batch_size", args.test_batch_size)
+        mlflow.log_param("model", args.model)
+        mlflow.log_param("feature_extract", args.feature_extract)
+        mlflow.log_param("use_pretrained", args.use_pretrained)
+
+        # Train and evaluate
+        history = train_and_evaluate(
+            model=model,
+            trainloader=trainloader,
+            testloader=testloader,
+            criterion=criterion,
+            optimizer=optimizer,
+            device=device,
+            epochs=args.epochs,
+            use_mlflow=True,
+        )
+        # Save the model
+        torch.save(model.state_dict(), f"pokemon_{args.model}.pth")
+        mlflow.log_artifact(f"pokemon_{args.model}.pth")
+        mlflow.end_run()