| 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 |
|
|
| |
| mask8 = (mask * 255).astype(np.uint8) |
|
|
| |
| contours, _ = cv2.findContours(mask8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) |
|
|
| yolo_segments = [] |
| for contour in contours: |
| if cv2.contourArea(contour) < 100: |
| continue |
|
|
| contour = contour.squeeze().astype(float) |
| if contour.ndim != 2: |
| continue |
|
|
| |
| 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 |
