Use the integer label from ONNX (#180)

ONNX models return labels as ints, but they're still useful for
filtering etc. Convert the int to a string and use it as the label.

This allows, for example, matching just breasts (2) or filtering out
exposed bellies (0) from NudeNet.

Co-authored-by: Sean Lynch <seanl@literati.org>
This commit is contained in:
Sean Lynch
2023-09-26 00:30:09 +09:00
committed by GitHub
co-authored by Sean Lynch
parent 10f2156e16
commit 6b8a4a37fc
+2 -2
View File
@@ -666,8 +666,8 @@ class ONNXDetector:
cropped_mask[y1 - crop_y1:y2 - crop_y1, x1 - crop_x1:x2 - crop_x1] = 1
cropped_mask = dilate_mask(cropped_mask, dilation)
# make items
item = SEG(None, cropped_mask, scores[i], crop_region, item_bbox, None, None)
# make items. just convert the integer label to a string
item = SEG(None, cropped_mask, scores[i], crop_region, item_bbox, str(labels[i]), None)
result.append(item)
shape = h, w