File size: 2,978 Bytes
8078d22 |
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 |
import numpy as np
from PIL import Image
import cv2
import os
import shutil
def blend_mask_with_image(image, mask, color):
"""Blend the mask with the original image using a transparent color overlay."""
mask_rgb = np.stack([mask * color[i] for i in range(3)], axis=-1)
blended = (0.7 * image + 0.3 * mask_rgb).astype(np.uint8)
return blended
def save_mask_as_png(mask, path):
"""Save the binary mask as a PNG."""
mask_image = Image.fromarray((mask * 255).astype(np.uint8))
mask_image.save(path)
def convert_mask_to_yolo(mask_path, image_path, class_id, output_path, append=False):
"""
Convert a binary mask to YOLO-compatible segmentation labels.
Args:
mask_path (str): Path to the binary mask image.
image_path (str): Path to the corresponding image.
class_id (int): Class ID (e.g., 0 for void, 1 for chip).
output_path (str): Path to save the YOLO label (.txt) file.
append (bool): Whether to append labels to the file.
Returns:
None
"""
try:
# Load the binary mask
mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
if mask is None:
raise ValueError(f"Mask not found or invalid: {mask_path}")
# Load the corresponding image to get dimensions
image = cv2.imread(image_path)
if image is None:
raise ValueError(f"Image not found or invalid: {image_path}")
h, w = image.shape[:2] # Image height and width
# Find contours in the mask
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
# Determine file mode: "w" for overwrite or "a" for append
file_mode = "a" if append else "w"
# Open the output .txt file
with open(output_path, file_mode) as label_file:
for contour in contours:
# Simplify the contour points to reduce the number of vertices
epsilon = 0.01 * cv2.arcLength(contour, True) # Tolerance for approximation
contour = cv2.approxPolyDP(contour, epsilon, True)
# Normalize contour points (polygon vertices)
normalized_vertices = []
for point in contour:
x, y = point[0] # Extract x, y from the point
x_normalized = x / w
y_normalized = y / h
normalized_vertices.extend([x_normalized, y_normalized])
# Write the polygon annotation to the label file
if len(normalized_vertices) >= 6: # At least 3 points required for a polygon
label_file.write(f"{class_id} " + " ".join(f"{v:.6f}" for v in normalized_vertices) + "\n")
print(f"YOLO segmentation label saved: {output_path}")
except Exception as e:
print(f"Error converting mask to YOLO format: {e}")
raise RuntimeError(f"Failed to convert {mask_path} for class {class_id}: {e}")
|