From 6b8a4a37fc6f528a67f8665c97d341e8bbe284d0 Mon Sep 17 00:00:00 2001 From: Sean Lynch <167674+seanlynch@users.noreply.github.com> Date: Mon, 25 Sep 2023 11:30:09 -0400 Subject: [PATCH] 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 --- modules/impact/core.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/modules/impact/core.py b/modules/impact/core.py index 9282963..2f1713d 100644 --- a/modules/impact/core.py +++ b/modules/impact/core.py @@ -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