This repository contains the official training and validation scripts for the Fishial.ai fish segmentation, detection, and classification models.
The easiest way to get started is by using our Google Colab Notebook. It allows you to run segmentation and classification models directly in your browser, using Google's cloud infrastructure or your own local machine.
To run the scripts on your own machine, follow these steps.
1. Clone the repository:
git clone <repository-url>
cd <repository-directory>2. Install dependencies: Make sure you have Python 3.x installed.
pip3 install -r requirements.txt3. Verify requirements:
python3 --versionshould report Python 3.10 or newer.- Use the pinned dependencies in
requirements.txtand recreate your virtual environment if you switch Python versions.
First, make the shell scripts executable:
chmod +x segmentation.sh classification.sh object_detection.sh./segmentation.sh -c <your_coco_file> -i <your_images_dir> -d <your_segmentation_ds_name> -s <your_save_dir>-c: Path to your COCO annotations file.-i: Path to the directory containing your images.-d: Name for your segmentation dataset.-s: Directory where the trained model will be saved.
./classification.sh -p <your_classif_images_dir> -i <your_classif_input_dir> -a <your_annotation_file> -n <your_classification_ds_name>-p: Path to the classification images directory.-i: Path to the classification input directory.-a: Path to your annotation file.-n: Name for your classification dataset.
./object_detection.sh -d <your_detection_ds> -o <your_yolo_output_dir> -n <num_classes> -y <data_yaml> -p <project_dir> -r <run_name>-d: Your detection dataset.-o: Output directory for YOLO results.-n: Number of classes.-y: Path to thedata.yamlfile.-p: Project directory.-r: Name for the specific run.
Here's a breakdown of the most important scripts and modules in this project.
train_scripts/classification/auto_train_cross.py: Automatically trains a classification model using the cross-entropy loss function. It saves the best-performing checkpoint based on validation accuracy.train_scripts/classification/auto_train_triplet.py: Trains a classification model using Triplet or Quadruplet loss. It saves the best checkpoint based on the k-metric on the validation set.train_scripts/segmentation/train.py: A basic script to train a segmentation model using the Detectron2 API.train_scripts/segmentation/train_copy_paste.py: Trains a segmentation model using Detectron2 with the "Copy-Paste" data augmentation technique.
helper/ExportModelToTorchscript.ipynb: A Jupyter Notebook to convert PyTorch classification models and Detectron2 segmentation models to the TorchScript format for optimized deployment.helper/classification/CreateDataBaseTensor.py: Generates an embedding tensor from a trained classification network. This tensor is used for efficient inference.helper/classification/CreateDatasetAndTrain.ipynb: A script to create training/testing datasets from a Fishial COCO export and subsequently train a network.
module/classification_package/src/model.py&.../utils.py: These files contain the core implementation of the classification pipeline.
The repository combines segmentation, detection, and classification tooling with helper notebooks and export utilities:
- Segmentation pipelines: Detectron2-based training scripts, augmentations, and export helpers.
- Object detection: YOLO training, validation, and inference helpers focused on bounding-box detection.
- Classification: Cross-entropy and metric-based training scripts plus inference helpers that produce class embeddings.
- Helpers & notebooks: Utility notebooks, embedding database tools, and TorchScript export utilities keep deployment workflows tidy.
The following checkpoints are organized by task. TorchScript-ready artifacts are highlighted when available.
| Model | Notes | Download |
|---|---|---|
| ⭐ YOLO v26 Fish Detector | TorchScript; latest export | download |
| YOLOv12 Medium Fish Detector | TorchScript; previous medium-sized detector | download |
| YOLOv10 Medium Fish Detector | TorchScript; earlier balanced model | download |
| Model | Notes | Download |
|---|---|---|
| ⭐ beitv2_base_patch16_224 (775 classes) | Embedding size 512; TorchScript export | download |
| beitv2_base_patch16_224 (640 classes) | Embedding 512; previous TorchScript pack | download |
| ConvNeXt Tiny (640 classes) | Embedding 256; TorchScript | download |
| ConvNeXt Tiny (426 classes) | Embedding 128; TorchScript | download |
| ResNet18 v6 model pack (289 classes) | ResNet-based classifier bundle | download |
| ResNet18 v5 model pack (184 classes) | Legacy ResNet18 set | download |
| ResNet18 v4 model pack (184 classes) | Early ResNet18 release | download |
| ResNet18 DataBase Tensor | Embedding tensor from train+test | download |
| ResNet18 Embedding 256 V2.0 | Full embedding checkpoint | download |
| ResNet18 Binary Classification | Binary ResNet18 checkpoint | download |
| ResNet18 Cross Entropy V1.0 | High-accuracy ResNet18 checkpoint | download |
| Model | Notes | Download |
|---|---|---|
| ⭐ FPN w/ ResNet18 (segmentation) | TorchScript-ready backbone (img size 416) | download |
| MaskRCNN Fish Segmentation (Updated 21.08.2023) | Detectron2 checkpoint | download |
| MaskRCNN Fish Segmentation (Updated 21.08.2023, TorchScript) | TorchScript export | download |
| MaskRCNN Fish Segmentation (Updated 15.11.2022) | Older Detectron2 checkpoint | download |
| MaskRCNN Fish Segmentation (Updated 29.06.2022) | Legacy checkpoint | download |
A JSON file containing the names of all fish classes recognized by the latest model can be found here: labels.json.
This project is licensed under the MIT License. See the LICENSE file for details.