Spaces:
Sleeping
Sleeping
File size: 1,984 Bytes
8c440cf | 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 | import io
import base64
import numpy as np
import cv2
from PIL import Image
def postprocess_mask(mask):
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
cleaned = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
cleaned = cv2.morphologyEx(cleaned, cv2.MORPH_CLOSE, kernel)
num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(cleaned, connectivity=8)
if num_labels <= 1:
return cleaned
largest_label = 1 + np.argmax(stats[1:, cv2.CC_STAT_AREA])
result = np.zeros_like(cleaned)
result[labels == largest_label] = 255
return result
def parse_mask_data(mask_data, h, w):
if isinstance(mask_data, str):
mask_bytes = base64.b64decode(mask_data)
mask_img = Image.open(io.BytesIO(mask_bytes))
mask_np = np.array(mask_img.convert('L'))
elif isinstance(mask_data, dict) and 'data' in mask_data:
mask_bytes = base64.b64decode(mask_data['data'])
mask_img = Image.open(io.BytesIO(mask_bytes))
mask_np = np.array(mask_img.convert('L'))
else:
mask_np = np.array(mask_data, dtype=np.uint8)
if mask_np.shape != (h, w):
resample = getattr(Image, 'LANCZOS', getattr(Image, 'NEAREST', None))
mask_pil = Image.fromarray(mask_np)
mask_pil = mask_pil.resize((w, h), resample)
mask_np = np.array(mask_pil)
binary_mask = np.zeros_like(mask_np, dtype=np.uint8)
binary_mask[mask_np > 127] = 255
return postprocess_mask(binary_mask)
def render_overlay(image, mask):
overlay = image.copy()
mask_bool = mask == 255
cyan_tint = np.zeros_like(image)
cyan_tint[:, :] = (255, 255, 0)
overlay[mask_bool] = cv2.addWeighted(image, 0.75, cyan_tint, 0.25, 0)[mask_bool]
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
cv2.drawContours(overlay, contours, -1, (0, 255, 255), 3)
_, buffer = cv2.imencode('.png', overlay)
return base64.b64encode(buffer).decode('utf-8')
|