46 lines
1.0 KiB
Python
46 lines
1.0 KiB
Python
import numpy as np
|
|
import cv2
|
|
from segment_anything import SamPredictor
|
|
|
|
|
|
def dilate_mask(mask, dilate_factor=15):
|
|
mask = mask.astype(np.uint8)
|
|
mask = cv2.dilate(
|
|
mask,
|
|
np.ones((dilate_factor, dilate_factor), np.uint8),
|
|
iterations=1
|
|
)
|
|
return mask
|
|
|
|
|
|
def predict_masks_with_sam(
|
|
img,
|
|
point_coords,
|
|
point_labels,
|
|
ckpt_p,
|
|
):
|
|
point_coords = np.array(point_coords)
|
|
point_labels = np.array(point_labels)
|
|
predictor = SamPredictor(ckpt_p)
|
|
|
|
predictor.set_image(img)
|
|
masks, scores, logits = predictor.predict(
|
|
point_coords=point_coords,
|
|
point_labels=point_labels,
|
|
multimask_output=True,
|
|
)
|
|
return masks, scores, logits
|
|
|
|
|
|
def get_masks(img, latest_coords, dilate_kernel_size, sam_ckpt):
|
|
masks, _, _ = predict_masks_with_sam(
|
|
img,
|
|
[latest_coords],
|
|
[1],
|
|
sam_ckpt,
|
|
)
|
|
masks = masks.astype(np.uint8) * 255
|
|
|
|
masks = [dilate_mask(mask, dilate_kernel_size) for mask in masks]
|
|
return masks
|