tyriaa/Segmentation_project

1import numpy as np2from PIL import Image3import cv24import os5import shutil6 7def blend_mask_with_image(image, mask, color):8    """Blend the mask with the original image using a transparent color overlay."""9    mask_rgb = np.stack([mask * color[i] for i in range(3)], axis=-1)10    blended = (0.7 * image + 0.3 * mask_rgb).astype(np.uint8)11    return blended12 13def save_mask_as_png(mask, path):14    """Save the binary mask as a PNG."""15    mask_image = Image.fromarray((mask * 255).astype(np.uint8))16    mask_image.save(path)17 18def convert_mask_to_yolo(mask_path, image_path, class_id, output_path, append=False):19    """20    Convert a binary mask to YOLO-compatible segmentation labels.21 22    Args:23        mask_path (str): Path to the binary mask image.24        image_path (str): Path to the corresponding image.25        class_id (int): Class ID (e.g., 0 for void, 1 for chip).26        output_path (str): Path to save the YOLO label (.txt) file.27        append (bool): Whether to append labels to the file.28 29    Returns:30        None31    """32    try:33        # Load the binary mask34        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)35        if mask is None:36            raise ValueError(f"Mask not found or invalid: {mask_path}")37 38        # Load the corresponding image to get dimensions39        image = cv2.imread(image_path)40        if image is None:41            raise ValueError(f"Image not found or invalid: {image_path}")42 43        h, w = image.shape[:2]  # Image height and width44 45        # Find contours in the mask46        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)47 48        # Determine file mode: "w" for overwrite or "a" for append49        file_mode = "a" if append else "w"50 51        # Open the output .txt file52        with open(output_path, file_mode) as label_file:53            for contour in contours:54                # Simplify the contour points to reduce the number of vertices55                epsilon = 0.01 * cv2.arcLength(contour, True)  # Tolerance for approximation56                contour = cv2.approxPolyDP(contour, epsilon, True)57 58                # Normalize contour points (polygon vertices)59                normalized_vertices = []60                for point in contour:61                    x, y = point[0]  # Extract x, y from the point62                    x_normalized = x / w63                    y_normalized = y / h64                    normalized_vertices.extend([x_normalized, y_normalized])65 66                # Write the polygon annotation to the label file67                if len(normalized_vertices) >= 6:  # At least 3 points required for a polygon68                    label_file.write(f"{class_id} " + " ".join(f"{v:.6f}" for v in normalized_vertices) + "\n")69 70        print(f"YOLO segmentation label saved: {output_path}")71 72    except Exception as e:73        print(f"Error converting mask to YOLO format: {e}")74        raise RuntimeError(f"Failed to convert {mask_path} for class {class_id}: {e}")75