File size: 2,703 Bytes
bb14d6a | 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 | import numpy as np
import cv2
import pycocotools.mask as mask_utils
from PIL import ImageDraw
def compute_iou(mask_a, mask_b):
intersection = np.logical_and(mask_a, mask_b).sum()
union = np.logical_or(mask_a, mask_b).sum()
return 0.0 if union == 0 else intersection / union
def mask_to_bbox(mask):
ys, xs = np.where(mask)
if len(xs) == 0:
return None
return xs.min(), ys.min(), xs.max(), ys.max()
def mask_to_rle(mask, json_safe=True):
rle = mask_utils.encode(np.asfortranarray(mask.astype(np.uint8)))
if json_safe:
rle["counts"] = rle["counts"].decode("ascii")
return rle
def rle_to_mask(rle):
rle = rle.copy()
if isinstance(rle["counts"], str):
rle["counts"] = rle["counts"].encode("ascii")
return mask_utils.decode(rle)
def uncompressed_rle_to_mask(rle):
"""Decode COCO-style uncompressed RLE into a binary mask (0/1)."""
h, w = rle["size"]
counts = rle["counts"]
mask = np.zeros(h * w, dtype=np.uint8)
val = 0
idx = 0
for c in counts:
mask[idx:idx + c] = val
idx += c
val = 1 - val
mask = mask.reshape((h, w), order='F')
return mask
def mask_to_yolo(mask, class_id=0):
"""Convert a binary mask (0/1) into YOLO polygon segmentation format."""
h, w = mask.shape
# ensure 8-bit binary mask
mask8 = (mask * 255).astype(np.uint8)
# find outer contours only
contours, _ = cv2.findContours(mask8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
yolo_segments = []
for contour in contours:
if cv2.contourArea(contour) < 100: # ignore tiny noise
continue
contour = contour.squeeze().astype(float)
if contour.ndim != 2:
continue
# normalize to [0,1]
contour[:, 0] = contour[:, 0] / float(w)
contour[:, 1] = contour[:, 1] / float(h)
coords = contour.flatten().tolist()
yolo_segments.append(f"{class_id} " + " ".join(f"{x:.6f}" for x in coords))
return yolo_segments
def rle_to_yolo(rle, class_id=0):
mask = rle_to_mask(rle)
return mask_to_yolo(mask, class_id)
def uncompressed_rle_to_yolo(rle, class_id=0):
mask = uncompressed_rle_to_mask(rle)
return mask_to_yolo(mask, class_id)
def draw_yolo_on_pil(image, yolo_segments, color=(0,255,0)):
img = image.convert("RGB")
draw = ImageDraw.Draw(img)
w, h = img.size
for seg in yolo_segments:
parts = seg.strip().split()
class_id = int(parts[0])
coords = np.array([float(x) for x in parts[1:]]).reshape(-1, 2)
points = [(x * w, y * h) for x, y in coords]
draw.line(points + [points[0]], fill=color, width=2)
return img |