oysters.py

"""
Python3 script for YOLO prediction and sizing of  oyster objects.

 To ensure compatible dependencies, set up a new venv virtual environment
 and install  all required dependencies with just one installation after updating
  the PIP installer.
 $ pip install -U pip
 $ pip install ultralytics
"""

from ultralytics import YOLO

results_dir = 'oyster_yolo_results'
predict_name = 'example_oyster_results'
model = YOLO( 'aquaculture_oysters.pt')
oyster_img =  'oyster_sample.jpg'

# The size standard disk diameter. In this case, as millimeters, 
#  using a cutout from a 3 inch hockey puck.
std_mm = 76.2  

# Run inference with these arguments
prediction_kwargs = dict(
    save=True, # True Saves image to project/name.
    save_txt=True,  # Relative bounding box coordinates. up to conf limit.
    save_conf=True,  # Adds conf scores to bbox coordinates
    augment=False,
    show_labels=True,  # True shows class labels in the saved image.
    line_width=2, 
    project=results_dir,   # bbox txt goes into runs/detect/labels.
    name=predict_name,
    exist_ok=False,  # True, overwrites existing project/name
)

results = model.predict(
    source=oyster_img,
    half=False,  # False is much faster than True
    imgsz=960,
    conf=0.7,
    device='cpu',
    iou=0.7,
    max_det=400,
    classes=[0,1],  #  0 is disks, 1 is oysters.
    **prediction_kwargs
)

# Display detection bounding boxes with labels and confidence values.
results[0].show(line_width=2)

# Calculate and print size metrics, Are based on longest bounding box dimension.
# Does not account for longer  lengths of oysters oriented on the diagonal.
#  An empirically determined correction factor is recommended for better accuracy.
#  A good correction function is used in the companion GitHub repository.
boxes_tensor = results[0].boxes.xywh
boxes_numpy = boxes_tensor.cpu().numpy().astype(int)
w_h = boxes_numpy[:, 2:4]
px_lengths = w_h.max(axis=1)  

classes = results[0].boxes.cls
class_distribution = classes.cpu().numpy().astype(int)
num_disks = len(class_distribution[class_distribution == 0])
num_oysters = len(class_distribution[class_distribution == 1])
oyster_px_sizes = px_lengths[class_distribution == 1]
disk_px_sizes = px_lengths[class_distribution == 0]
px2mm_factor = std_mm / disk_px_sizes.mean()
oyster_mean_mm = round(oyster_px_sizes.mean() * px2mm_factor, 2)
smallest = round(oyster_px_sizes.min() * px2mm_factor, 2)
largest = round(oyster_px_sizes.max() * px2mm_factor, 2)

print(f'Mean oyster size, mm: {oyster_mean_mm}')
print(f'Oyster size range, mm: [{smallest} -- {largest}]')