feat: only process sam_options.dino_erode_or_dilate when set

This commit is contained in:
Manuel Schmid
2024-06-10 23:47:09 +02:00
parent ab3f4df935
commit b5b33361c9
+7 -6
View File
@@ -83,12 +83,13 @@ def generate_mask_from_image(image: np.ndarray, mask_model: str = 'sam', extras=
if boxes.size(0) > 0: if boxes.size(0) > 0:
sam_predictor.set_image(image) sam_predictor.set_image(image)
for index in range(boxes.size(0)): if sam_options.dino_erode_or_dilate != 0:
assert boxes.size(1) == 4 for index in range(boxes.size(0)):
boxes[index][0] -= sam_options.dino_erode_or_dilate assert boxes.size(1) == 4
boxes[index][1] -= sam_options.dino_erode_or_dilate boxes[index][0] -= sam_options.dino_erode_or_dilate
boxes[index][2] += sam_options.dino_erode_or_dilate boxes[index][1] -= sam_options.dino_erode_or_dilate
boxes[index][3] += sam_options.dino_erode_or_dilate boxes[index][2] += sam_options.dino_erode_or_dilate
boxes[index][3] += sam_options.dino_erode_or_dilate
if sam_options.dino_debug: if sam_options.dino_debug:
from PIL import ImageDraw, Image from PIL import ImageDraw, Image