# Shark Classification with Vision Transformers

This repository contains code for training and evaluating Vision Transformer (ViT) models for shark classification. The models perform binary shark identification (`train_SI_vit.py`), genus classification (`train_gen_vit.py`) and genus-specific species classification (`train_spec_vit.py`). Note that the binary identification and genus models both have one model each, while the species model has one model for each **data-sufficient** genus.

## Table of Contents

1. [Usage](#usage)
   - [Installation](#installation)
   - [Running the code](#running-the-code)
   - [Example](#example)
2. [Future Work](#future-work)
   - [Data pipeline](#data-pipeline)
   - [Docker](#docker)
   - [cron-job](#cron-job)
3. [More Info](#more-info)

---

## Usage

### Installation

To install the required packages, run:

```
pip install -r requirements.txt
```

This project used Python 3.10.10 for development.

### Running the code

The main entry point for ViT training is `run_training.py`. It accepts command-line arguments to specify which models to train and an optional output path to redirect output files.

```
python run_training.py -m {SI, gen, spec} [-o OUTPUT_PATH]
```

- `-m`, `--models`: Specify the models to train. Choose from 'SI' (Binary Classifier), 'gen' (Genus Classifier), and 'spec' (Species Classifier). You can train multiple models by providing a space-separated list, e.g.,`-m SI gen`.

- `-o`, `--outpath`: (Optional) Specify the output path for the trained models and other outputs. Defaults to `./out`.

### Example

To train the Binary Classification model and the Genus model, and save the output in the ./trained_models directory:

```
python run_training.py -m SI gen -o ./trained_models
```

This will train the specified models using the default parameters provided in `run_training.py`. You can modify the parameters (e.g., image size, batch size, number of epochs) directly in the script if needed.

After training, the trained models and their corresponding logs will be saved in the specified output directory (`./trained_models` in this example).

---

## Future Work

At the time of completion of this project, there remains some issues to address:

- (URGENT) Figure out why accuracy is so low in refactored code. It appears to be a trivial bug that has to do with `flow_from_dataframe` in `data_utils.py`
- Develop more standardized validation monitor outputs in order to scale pipeline.
- Containerize the training using `./Dockerfile` and set up a pipeline to have the trained models be deployed on the [SharkPulse](http://sharkpulse.cnre.vt.edu/) website
- Automate training via a cronjob

### Data pipeline

Namely, the images stored across different folders in the `sp2.cs.vt.edu` node have some overlap.

Each table in the `pelagic` PSQL database indexes files based on their filename without a file extension or path; This presents an issue for indexing the absolute paths of images (a task necessary for fetching the training input) as they are stored **non-disjointly** in various folders `~/sd_images`, `~/instaval_images`, and others.


This repository includes the middleware mapping of tables to file location - `data_utils.aggregate_tables()` - but it does not scale easily.

Additionally, the databases should all share common attribute names so that middleware can easily interface with them (eg. all file names should be stored as `img_name`). Imposing this order can be done manually by editing the validation monitor software OR via Object-Oriented-Programming a type `<Scraper>` which has a defined schema for constructing tables.

### Docker

To build a new image (do this after modifying any code), run (in the current directory):
`docker build -t sp-train-vit .`

To start training, run:
`docker run -v ~/SDv3/saved_models:/app/saved_models -v ~/sd_images:/app/sd_images sp-train-vit`

Note that the `-v` arg mounts volumes for access to the filesystem.
Anything not mounted is ephemeral and will be discarded after container stops.
This is necessary to have access to the various image folders (eg. `sd_images`.)

Note this implies that in order to run on another machine you MUST transfer the necessary image folders (See SFTP or SCP).

Containerized machine learning relies on the NVIDIA Container Toolkit; it enables containers to use CUDA cores.

### `cron`-job

Automated monthly training can be implemented by adding the following string to your server's `crontab`:

```
0 0 1 * * /usr/bin/python /path/to/your/script/run_training.py
```

NOTE: it would be better to implement a script `check_data_influx.sh` that sees if the training data has grown sufficiently since the last run.

---

## More Info

To read more about docker, visit:
`https://www.digitalocean.com/community/tutorials/working-with-docker-containers`

Keras `ImageDataGenerator` docs (contains info on building `training_imgaes` from directory vs dataframe):
`https://www.tensorflow.org/api_docs/python/tf/keras/preprocessing/image/ImageDataGenerator#flow_from_dataframe`

SharkPulse website:
`http://sharkpulse.cnre.vt.edu/`

Paper for the original SharkPulse model:
`https://www.sciencedirect.com/science/article/abs/pii/S1574954122001236`

Paper for the Vision Transformer model:
`INSERT LINK HERE`
