Fathomnet train
This notebook by MBARI demonstrates how to train a YOLOv5 object detection model using freely available FathomNet data using AWS. See the full documentation on how to setup the deepsea-ai module or check github for questions.
Setup¶
If you have successfully setup your Anaconda environment, you should see the kernel deepsea-ai-notebooks. Select that kernel.
A few housekeeping items:
- The packages installed require python3.7 or greater, so let's check the python version first, then install the additional packages used in this notebook
- Verify your AWS account is working by listing your buckets
!python --version
!pip install -U -q fathomnet pillow
!pip install -U -q deepsea-ai
!aws s3 ls
Select concepts¶
Let's choose some concepts that are in the FathomNet database and put them in a list. We are also going to put them into a file, one line per concept in alphabetical order which will be used during training.
concepts = ["Rathbunaster californicus", "Holothuroidea", "Strongylocentrotus fragilis"]
with open('mynames.txt', 'w+') as n:
for c in sorted(concepts):
n.write(c + '\r\n')
#This should be in alphabetical order
!cat mynames.txt
Let's call the fathomnet.api.images.find_by_concept
function to query how many images are available for these concepts.
from fathomnet.api import images
for c in concepts:
available_images = images.find_by_concept(c)
print(f'{c} images: {len(available_images)}')
Download the images and format the labels¶
This code will:
- Download the images and bounding box annotations from FathomNet and save them to the
trainingdata/
folder. - Format the bounding boxes into a format compatible with the object detection model. The YOLOv5 model uses a very simple format which is a text file, one line per annotation, e.g.
0 101 33 83 199
which corresponds to```<label index 0-based>
import requests
from pathlib import Path
from progressbar import progressbar
from fathomnet.api import images
from PIL import Image
from fathomnet.models import GeoImageConstraints
# Create a directory for the images and labels
data_dir = Path.cwd() / 'trainingdata'
image_dir = data_dir / 'images'
image_dir.mkdir(exist_ok=True, parents=True)
label_dir = data_dir / 'labels'
label_dir.mkdir(exist_ok=True, parents=True)
# Download each image and create a label file for training
image_paths = []
for c in concepts:
# Constrain to only 50 images
concept_constrained = GeoImageConstraints(concept=c, limit=50)
concept_images = images.find(concept_constrained)
print(f'Downloading {c} images: {len(concept_images)}')
for image in progressbar(concept_images):
# Export labels with the same name as the unique identifier appended with .txt per darknet format
label = label_dir / f'{image.uuid}.txt'
print(image.boundingBoxes)
# Export to darknet format, which is one line per annotation <label index 1-based> <x> <y> <width> <height>
with label.open("w+") as l:
for b in image.boundingBoxes:
# Only save concepts in our list
if b.concept in concepts:
l.write(f'{concepts.index(b.concept)} {b.x/image.width} {b.y/image.height} {b.width/image.width} {b.height/image.height}\n')
# # Format our image file name as the image UUID + .jpg
image_path = image_dir / f'{image.uuid}.jpg'
image_paths.append(image_path)
if image_path.exists(): # Skip re-downloading images
continue
# # Download the image
image_raw = requests.get(image.url, stream=True).raw
pil_image = Image.open(image_raw)
# Convert to RGB (ensures consistent colorspace)
pil_image = pil_image.convert('RGB')
# Save the image
pil_image.save(image_path)
Train¶
Now that we have some training data, let's train it for a few cycles.
Before we train it, we need to split the data into training, testing with the deepsea-ai split command. This will split the data, then compressed it into images.tar.gz and labels.tar.gz files to reduce the data size.
The training data downloaded should be organized something like:
├── trainingdata
│ │ ├── images
│ │ │ └── image1.jpg
│ │ │ └── image2.jpg
│ │ ├── labels
│ │ │ └── image1.txt
│ │ │ └── image2.txt
and after running the split command, you should have just two files
├── splitdata
│ │ ├── images.tar.gz
│ │ ├── labels.tar.gz
!deepsea-ai split -i trainingdata -o splitdata
Train the object detection model for a few epochs¶
The object detector model is the highly performant YOLOv5 model. This represents a good tradeoff between cost and performance. See training for the full documentation.
To track cost, we will setup unique tags for our project, and pass that to the --config
option.
%%writefile deepseailab.txt
[docker]
yolov5_container = mbari/deepsea-yolov5:1.1.2
strongsort_container = mbari/strongsort-yolov5:5c05df0
[aws]
account_id = 548531997526
sagemaker_arn = arn:aws:iam::548531997526:role/DeepSeaAI
model = s3://deepsea-ai-548531997526-models/yolov5x_mbay_benthic_model.tar.gz
track_config = s3://deepsea-ai-548531997526-track-conf/strong_sort_benthic.yaml
videos = s3://deepsea-ai-548531997526-videos
models = s3://deepsea-ai-548531997526-models
tracks = s3://deepsea-ai-548531997526-tracks
[aws_public]
model = s3://902005-public/models/yolov5x_mbay_benthic_model.tar.gz
track_config = s3://902005-public/models/track-config/strong_sort_benthic.yaml
video_ex1 = s3://902005-public/video/1sec/V4361_20211006T162656Z_h265_1sec.mp4
video_ex2 = s3://902005-public/video/1sec/V4361_20211006T163256Z_h265_1sec.mp4
video_ex3 = s3://902005-public/video/1sec/V4361_20211006T163856Z_h265_1sec.mp4
[database]
site = http://localhost:4000
gql = %(site)/graphql
[tags]
organization = deepseailab
project_number = NA
stage = test
application = detection
Create a unique bucket name¶
Please note that bucket names need to be globally unique. Let's set the bucket name here with a randomly generated string to help with that. You can remove this and replace with your unique bucket name.
import uuid
uuid_short = str(uuid.uuid4())[:8]
uuid_short
Train¶
We set the number of epochs to 2 which is far too small to be useful, but we will test it first to make sure your AWS account works ok. We will use the ml.p3.2xlarge instance which is at the time of this writing $3.06 USD per hour. Once this is done, you might try it again for a longer period to see if the performance improves, or try a larger instance type, or try downloading more images. At the time of this writing, typically less than 40 epochs are needed; beyond that performance does not improve.
!deepsea-ai train --model yolov5x --instance-type ml.p3.2xlarge \
--labels splitdata/labels.tar.gz \
--images splitdata/images.tar.gz \
--label-map mynames.txt \
--config deepseailab.txt \
--input-s3 f's3://{uuid_short}-deepseailab-benthic-training/' \
--output-s3 f's3://{uuid_short}-deepseailab-benthic-checkpoints/' \
--epochs 2 \
--batch-size 2