Spaces:
Running
Running
File size: 3,355 Bytes
123489f |
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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
from PIL import Image
import numpy as np
from .base_segmenter import BaseSegmenter
from .painter import mask_painter, point_painter
mask_color = 3
mask_alpha = 0.7
contour_color = 1
contour_width = 5
point_color_ne = 8
point_color_ps = 50
point_alpha = 0.9
point_radius = 15
contour_color = 2
contour_width = 5
class SamControler:
def __init__(self, sam_pt_checkpoint, sam_onnx_checkpoint, model_type, device):
"""
initialize sam controler
"""
self.sam_controler = BaseSegmenter(sam_pt_checkpoint, sam_onnx_checkpoint, model_type, device)
self.onnx = model_type == "vit_t"
def first_frame_click(
self,
image: np.ndarray,
points: np.ndarray,
labels: np.ndarray,
multimask=True,
mask_color=3,
):
"""
it is used in first frame in video
return: mask, logit, painted image(mask+point)
"""
# self.sam_controler.set_image(image)
neg_flag = labels[-1]
if self.onnx:
onnx_coord = np.concatenate([points, np.array([[0.0, 0.0]])], axis=0)[None, :, :]
onnx_label = np.concatenate([labels, np.array([-1])], axis=0)[None, :].astype(np.float32)
onnx_coord = self.sam_controler.predictor.transform.apply_coords(onnx_coord, image.shape[:2]).astype(np.float32)
prompts = {
"point_coords": onnx_coord,
"point_labels": onnx_label,
"orig_im_size": np.array(image.shape[:2], dtype=np.float32),
}
else:
prompts = {
"point_coords": points,
"point_labels": labels,
}
if neg_flag == 1:
# find positive
masks, scores, logits = self.sam_controler.predict(
prompts, "point", multimask
)
mask, logit = masks[np.argmax(scores)], logits[np.argmax(scores), :, :]
prompts["mask_input"] = np.expand_dims(logit[None, :, :], 0)
masks, scores, logits = self.sam_controler.predict(
prompts, "both", multimask
)
mask, logit = masks[np.argmax(scores)], logits[np.argmax(scores), :, :]
else:
# find neg
masks, scores, logits = self.sam_controler.predict(
prompts, "point", multimask
)
mask, logit = masks[np.argmax(scores)], logits[np.argmax(scores), :, :]
assert len(points) == len(labels)
painted_image = mask_painter(
image,
mask.astype("uint8"),
mask_color,
mask_alpha,
contour_color,
contour_width,
)
painted_image = point_painter(
painted_image,
np.squeeze(points[np.argwhere(labels > 0)], axis=1),
point_color_ne,
point_alpha,
point_radius,
contour_color,
contour_width,
)
painted_image = point_painter(
painted_image,
np.squeeze(points[np.argwhere(labels < 1)], axis=1),
point_color_ps,
point_alpha,
point_radius,
contour_color,
contour_width,
)
painted_image = Image.fromarray(painted_image)
return mask, logit, painted_image
|