File size: 6,494 Bytes
bca2f18 | 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 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | """
SAM implementation using HuggingFace transformers for PyTorch 2.7 compatibility.
"""
import numpy as np
import torch
from PIL import Image
from typing import Union
# Import utilities
from ..util import HWC3, common_input_validate, resize_image_with_pad
class SamDetector:
def __init__(self, model_name="facebook/sam-vit-base"):
from transformers import SamModel, SamProcessor
self.model_name = model_name
self.processor = SamProcessor.from_pretrained(model_name)
self.model = SamModel.from_pretrained(model_name)
self.device = "cpu"
@classmethod
def from_pretrained(cls, pretrained_model_or_path=None, model_type="vit_t", filename="mobile_sam.pt", subfolder=None):
model_mapping = {
"vit_t": "facebook/sam-vit-base",
"vit_b": "facebook/sam-vit-base",
"vit_l": "facebook/sam-vit-large",
"vit_h": "facebook/sam-vit-huge"
}
if filename and isinstance(filename, str):
if "mobile_sam" in filename.lower():
model_name = "facebook/sam-vit-base"
elif "sam_vit_h" in filename.lower():
model_name = "facebook/sam-vit-huge"
elif "sam_vit_l" in filename.lower():
model_name = "facebook/sam-vit-large"
elif "sam_vit_b" in filename.lower():
model_name = "facebook/sam-vit-base"
else:
model_name = model_mapping.get(model_type, "facebook/sam-vit-base")
else:
model_name = model_mapping.get(model_type, "facebook/sam-vit-base")
return cls(model_name)
def to(self, device):
self.model = self.model.to(device)
self.device = device
return self
def generate_automatic_masks(self, input_image):
if isinstance(input_image, np.ndarray):
pil_image = Image.fromarray(input_image)
else:
pil_image = input_image
height, width = pil_image.size[1], pil_image.size[0]
points_per_side = max(8, min(24, width // 64, height // 64))
grid_points = []
for i in range(points_per_side):
for j in range(points_per_side):
x = int((j + 0.5) * width / points_per_side)
y = int((i + 0.5) * height / points_per_side)
x_offset = int((np.random.random() - 0.5) * (width / points_per_side * 0.3))
y_offset = int((np.random.random() - 0.5) * (height / points_per_side * 0.3))
x = max(5, min(width - 5, x + x_offset))
y = max(5, min(height - 5, y + y_offset))
grid_points.append([x, y])
batch_size = 16
all_masks = []
for i in range(0, len(grid_points), batch_size):
batch_points = grid_points[i:i + batch_size]
input_points = [batch_points]
inputs = self.processor(
images=pil_image,
input_points=input_points,
return_tensors="pt"
).to(self.device)
with torch.no_grad():
outputs = self.model(**inputs)
masks = self.processor.post_process_masks(
outputs.pred_masks,
inputs["original_sizes"],
inputs["reshaped_input_sizes"]
)[0]
masks_np = masks.cpu().numpy()
for j, mask in enumerate(masks_np):
mask_2d = mask[0] if len(mask.shape) > 2 else mask
area = int(mask_2d.sum())
if area > 100:
cleaned_mask = self._postprocess_mask(mask_2d)
cleaned_area = int(cleaned_mask.sum())
mask_dict = {
'segmentation': cleaned_mask,
'area': cleaned_area,
'stability_score': 0.88,
'point_coords': batch_points[j % len(batch_points)]
}
all_masks.append(mask_dict)
return all_masks
def _postprocess_mask(self, mask, min_region_area=100):
from scipy import ndimage
from skimage import morphology
binary_mask = mask.astype(bool)
filled_mask = ndimage.binary_fill_holes(binary_mask)
if filled_mask.any():
kernel_close = morphology.disk(5)
kernel_open = morphology.disk(3)
smoothed_mask = morphology.binary_closing(filled_mask, kernel_close)
smoothed_mask = morphology.binary_opening(smoothed_mask, kernel_open)
smoothed_mask = morphology.binary_closing(smoothed_mask, kernel_close)
else:
smoothed_mask = filled_mask
return smoothed_mask.astype(mask.dtype)
def show_anns(self, anns):
if len(anns) == 0:
return None
sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
h, w = anns[0]['segmentation'].shape
final_img = Image.fromarray(np.zeros((h, w, 3), dtype=np.uint8), mode="RGB")
for ann in sorted_anns:
m = ann['segmentation']
img = np.empty((m.shape[0], m.shape[1], 3), dtype=np.uint8)
for i in range(3):
img[:,:,i] = np.random.randint(255, dtype=np.uint8)
final_img.paste(Image.fromarray(img, mode="RGB"), (0, 0), Image.fromarray(np.uint8(m*255)))
return np.array(final_img, dtype=np.uint8)
def __call__(self, input_image: Union[np.ndarray, Image.Image]=None, detect_resolution=512, output_type="pil", upscale_method="INTER_CUBIC", **kwargs) -> Image.Image:
input_image, output_type = common_input_validate(input_image, output_type, **kwargs)
input_image, remove_pad = resize_image_with_pad(input_image, detect_resolution, upscale_method)
masks = self.generate_automatic_masks(input_image)
map = self.show_anns(masks)
if map is None:
map = np.zeros((input_image.shape[0], input_image.shape[1], 3), dtype=np.uint8)
detected_map = HWC3(remove_pad(map))
if output_type == "pil":
detected_map = Image.fromarray(detected_map)
return detected_map |