local_segment/utils/sam_utils.py

import numpy as np
from PIL import Image
import io
import base64

def show_mask(mask):
    h, w = mask.shape[-2:]
    mask = mask.astype(np.uint8)
    # 创建一个全黑的RGBA图像
    mask_image = np.zeros((h, w, 4), dtype=np.uint8)
    # 将掩码区域设为白色（前景）
    mask_image[mask > 0] = [0, 122, 204, 255]  # 白色前景
    mask_image[mask == 0] = [0, 0, 0, 0]       # 黑色背景
    return mask_image

def show_masks(masks, scores):
    base64_images = []
    for i, (mask, score) in enumerate(zip(masks, scores)):
        np_arr = show_mask(mask)
        arr_img = Image.fromarray(np_arr)
        
        # 将图像转换为base64
        buffered = io.BytesIO()
        arr_img.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()
        base64_images.append(img_str)
        
    return base64_images

def convert_to_serializable(obj):
    """递归将numpy和torch类型转换为可序列化的Python类型"""
    import torch
    if isinstance(obj, np.ndarray):
        return obj.tolist()  # 将numpy数组转换为Python列表
    elif isinstance(obj, torch.Tensor):
        return obj.detach().cpu().numpy().tolist()  # 将tensor转换为numpy再转为列表
    elif isinstance(obj, (np.floating, np.integer, np.bool_)):
        return obj.item()  # 将numpy标量转换为Python原生类型
    elif isinstance(obj, (list, tuple)):
        return [convert_to_serializable(item) for item in obj]
    elif isinstance(obj, dict):
        return {key: convert_to_serializable(value) for key, value in obj.items()}
    else:
        return obj