File size: 2,590 Bytes
afb9fe9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
"""
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}]')
|