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.txtFirst, 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.
We provide several pre-trained models for immediate use. The latest models are marked with ⭐.
| Model Description | Download Link |
|---|---|
| ⭐ Fish Detector BoundingBox - YOLOv12 Medium (img size 640, torchscript) | link |
| ⭐ Fish Classification - beitv2_base_patch16_224 (640 classes, embed 512, torchscript) | link |
| ⭐ Fish Classification - ConvNeXt Tiny (640 classes, embed 256, torchscript) | link |
| ⭐ Fish Segmentation - FPN w/ ResNet18 (img size 416, torchscript) | link |
| Fish Detector BoundingBox - YOLOv10 Medium (img size 640, torchscript) | link |
| Fish Classification - ConvNeXt Tiny (426 classes, embed 128, torchscript) | link |
| Fish Segmentation - FPN w/ ResNet18 (img size 416, torchscript) | link |
| MaskRCNN Fish Segmentation (Updated 21.08.2023) | link |
| MaskRCNN Fish Segmentation (Updated 21.08.2023, torchscript) | link |
| ResNet18 v6 model pack (289 classes) | link |
| ResNet18 v5 model pack (184 classes) | link |
| MaskRCNN Fish Segmentation (Updated 15.11.2022) | link |
| ResNet18 v4 model pack (184 classes) | link |
| ResNet18 DataBase Tensor | link |
| ResNet18 Fish Classification Embedding 256 V2.0 | link |
| ResNet18 Binary Classification | link |
| ResNet18 Fish Classification Cross Entropy V1.0 | link |
| MaskRCNN Fish Segmentation (Updated 29.06.2022) | link |
For more detailed results, see the Training README.
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.