You may want to check the accuracy of a segmentation model by its precision/recall metrics at the image level. The solution we present here will take the segmentation with most pixels. The variation could be:
Convert to OK/NG
Use a defect order to rank the labels instead of pixel size.
Content of custom/segmentation_to_classification.py.
from landinglens.model_iteration.sdk import BaseTransform, DataItem import numpy as np
class SegmentationToClassification(BaseTransform): """Transforms a segmentation output into a classification output. If there are NG pixels, the output will be the NG class with the most pixels; otherwise, it will be OK. """
def__init__(self, **params): """Any parameters defined in yaml file will be passed to init. Store the value passed in self so you can access them in the call method. """
def __call__(self, inputs: DataItem) -> DataItem: """Return a new DataItem with transformed attributes. DataItem has following attributes: image - input image. mask_scores, mask_labels - segmentation mask probabilities and classes. Returns ------- A named tuple class DataItem with the modified attributes. """ # Get scores and labels and raise error if not defined mask = inputs.mask_scores
if mask is None: raise TypeError("'mask_scores' not defined in inputs") output_depth = mask.shape[-1] labels = np.argmax(mask, -1)
# Count how many labels pixel are in total and sort them values, counts = np.unique(labels, return_counts=True) arg_sort = np.argsort(counts)[::-1] values_sorted = values[arg_sort] counts_sorted = counts[arg_sort]
# Get label and score (proportion) position = 0if len(values) == 1 else 1 label = int(values_sorted[position]) score = float(counts_sorted[position] / np.sum(counts_sorted))