Spaces:
Runtime error
Runtime error
File size: 3,021 Bytes
0dd537b |
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 |
from typing import IO, List
import torch
from segment_anything import SamPredictor, sam_model_registry, SamAutomaticMaskGenerator
from PIL import Image
import numpy as np
import io
def to_file(item) -> IO[bytes]:
# Create a BytesIO object
file_obj = io.BytesIO()
if isinstance(item, Image.Image):
item.save(file_obj, format='PNG')
if isinstance(item, np.ndarray):
np.save(file_obj, item)
# Reset the file object's position to the beginning
file_obj.seek(0)
# Return the file object
return file_obj
def get_sam(model_type, checkpoint_path, device=None):
if device is None:
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
sam = sam_model_registry[model_type](checkpoint=checkpoint_path)
sam.to(device=device)
return sam
def draw_mask(img: Image.Image, boolean_mask: np.ndarray, color: tuple, mask_alpha: float) -> Image.Image:
int_alpha = int(mask_alpha*255)
color_mask = Image.new('RGBA', img.size, color=color)
color_mask.putalpha(Image.fromarray(boolean_mask.astype(np.uint8)*int_alpha, mode='L'))
result = Image.alpha_composite(img, color_mask)
return result
def random_color():
return tuple(np.random.randint(0,255, 3))
def draw_masks(img: Image.Image, boolean_masks: np.ndarray) -> Image.Image:
img = img.copy()
for boolean_mask in boolean_masks:
img = draw_mask(img, boolean_mask, random_color(), 0.2)
return img
def cutout(img: Image.Image, boolean_mask: np.ndarray):
rgba_img = img.convert('RGBA')
mask = Image.fromarray(boolean_mask).convert("L")
rgba_img.putalpha(mask)
return rgba_img
def predict_conditioned(sam, pil_img, **kwargs):
rgb_arr = pil_image_to_rgb_array(pil_img)
predictor = SamPredictor(sam)
predictor.set_image(rgb_arr)
masks, quality, _ = predictor.predict(**kwargs)
return masks, quality
def predict_all(sam, pil_img):
rgb_arr = pil_image_to_rgb_array(pil_img)
mask_generator = SamAutomaticMaskGenerator(sam)
results = mask_generator.generate(rgb_arr)
masks = []
quality = []
for result in results:
masks.append(result['segmentation'])
quality.append(result['stability_score'])
masks = np.array(masks)
quality = np.array(quality)
return masks, quality
def pil_image_to_rgb_array(image):
if image.mode == "RGBA":
rgb_image = Image.new("RGB", image.size, (255, 255, 255))
rgb_image.paste(image, mask=image.split()[3]) # Apply alpha channel as the mask
rgb_array = np.array(rgb_image)
else:
rgb_array = np.array(image.convert("RGB"))
return rgb_array
def box_pts_to_xyxy(pt1, pt2):
"""convert box from pts format to XYXY
Args:
pt1 : (x1, y1) first corner of a box
pt2 : (x2, y2) second corner, diagonal to pt1
Returns:
xyxy: (x_min, y_min, x_max, y_max)
"""
x1, y1 = pt1
x2, y2 = pt2
return (min(x1, x2), min(y1, y2), max(x1, x2), max(y1, y2)) |