This project includes training and validation scripts for the fish segmentation and classification model.
Project website: www.fishial.ai
Install the dependencies.
$ pip3 install -r requirements.txt-
Runner.ipynb This jupyter notebook allows you to run segmentation and classification neural networks on Google Cloud or your computer, after downloading the files from the links below.
-
auto_train_cross.py is the script performs training automatically with different parameters the selected model using the cross entropy loss function. The checkpoint with the best performance on the validation dataset is saved to the output folder.
-
auto_train_triplet.py is the script performs training automatically with different the selected model using the (Triplet Quadruplet) loss function. The checkpoint with the best performance on the validation dataset according k-metric is saved to the output folder.
-
auto_train_cross.py is the script cut your dataset to specific maximum and minimum count images per class to the arbitrary way.
-
train.py is the basic set up script to train segmentation model using Detectrin2 API
-
train_copy_paste.py is the basic set up script to train segmentation model using Detectrin2 API with Copy Paste Augumentation.
-
ExportModelToTorchscript.ipynb This jupyter notebook allows convert Classification pytorch model to TorchScript format, and Detectron2 to Torchscript model.
-
(model.py, utils.py): These files contain the main classification pipline implementation.
-
CreateDataBaseTensor.py: This script is designed to get the attachment tensor for a trained neural network, the resulting tensor has dimensions {number of classes * maximum number of images for one class * embedding dimension} for classes in which the number of attachments is less than in the maximum class, a tensor with a value of 100 is added in order not to affect for inference
-
CreateDatasetAndTrain.py: This script allows you to create a training and test data set from the exported fishial coco file and train the neural network
| Model | link |
|---|---|
| MaskRCNN Fish Segmentation (Update 29.06.2022) | link |
| ResNet18 Fish Classification Cross Entropy V1.0 | link |
| ResNet18 Binary Classification | link |
| ResNet18 Fish Classification Embedding 256 V2.0 | link |
| ResNet18 DataBase Tensor | link |
| ResNet18 v4 model pack 184 classes | link |
| MaskRCNN Fish Segmentation (Update 15.11.2022) | link |
| ResNet18 v5 model pack 184 classes | link |
| ResNet18 v6 model pack 289 classes (latest) | link |
| MaskRCNN Fish Segmentation (Update 21.08.2023) (latest) | link |
| MaskRCNN Fish Segmentation (Update 21.08.2023) (latest) torchscript | link |
| in the folder with the classification model (V7) you will find an example of a python script (runner.py) how to run the model on your images |
Segmentation model has validated by mAP metric.
**MaskRCNN **
| AP | AP50 | AP75 | APs | APm | APl |
|---|---|---|---|---|---|
| 82.504 | 96.742 | 94.727 | 13.283 | 58.029 | 84.540 |
Classification model
Json file with the names of fish classes that the latest model recognizes can be found here: (labels)
