Spaces:
Runtime error
Runtime error
| 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)) |