label_schema.py · constructelligence/painting-vision-robotics-kit at main

painting-vision-robotics-kit / label_schema.py

Upload label_schema.py with huggingface_hub

2004e55 verified 29 days ago

2.92 kB

"""Shared lossless mask taxonomy for the painting vision dataset."""
import numpy as np

CLASSES = ("other", "wall_unpainted", "wall_painted", "wall_uncertain", "skirting",
           "switch_outlet", "ac_unit", "door", "window", "wall_obstacle")
IGNORE = 255
UNKNOWN = 254

# Visible surface material of wall pixels, predicted by the material head
# independently of paint state. Wallpaper is the visible covering, so a
# wallpapered gypsum wall is ``wallpaper``, not ``drywall``. Older datasets that
# only use 0/1 stay valid; 255 marks unknown material or non-wall pixels.
MATERIAL_CLASSES = ("other_wall_material", "drywall", "wallpaper")
LEGACY_MATERIAL_CLASSES = MATERIAL_CLASSES[:2]
# A smooth finished coat over board or plaster: the material is genuinely one of
# the two and the photo cannot tell which. Trained toward 50/50 drywall/other, so
# the model's honest answer there is low confidence (predict.py: "unknown").
MATERIAL_DRYWALL_OR_OTHER = 254
MATERIAL_VALUES = tuple(range(len(MATERIAL_CLASSES))) + (MATERIAL_DRYWALL_OR_OTHER, IGNORE)

# Surface condition of decided wall pixels, from a third head: what has to be
# prepared before paint goes on. ``paint_failure`` covers peeling, flaking,
# blistering and cracked coatings; ``rust`` covers corrosion and rust staining.
# Masks live in condition_masks/<split>/<stem>.png; 255 = unknown or not wall.
# Datasets without condition masks are simply unsupervised for this head.
CONDITION_CLASSES = ("sound", "paint_failure", "rust")
CONDITION_VALUES = tuple(range(len(CONDITION_CLASSES))) + (IGNORE,)

# Detection classes for the equipment around the wall: what has to be on site to
# do the job, and what must not be painted over. These are *objects*, not pixel
# labels - they never appear in the ten-class mask, so ``CLASSES`` above (and any
# trained checkpoint) is unaffected by adding or reordering them. ``site_objects.py``
# detects them without a model, from shape and colour cues.
OBJECT_CLASSES = ("paint_can", "paint_tray", "brush", "roller", "ladder", "drop_sheet")

# RGB values for lossless color-coded PNG masks.
RGB_PALETTE = {
    (0, 0, 0): 0,
    (0, 0, 255): 1,
    (0, 255, 0): 2,
    (255, 255, 0): 3,
    (255, 0, 255): 4,
    (255, 0, 0): 5,
    (0, 255, 255): 6,
    (255, 128, 0): 7,
    (128, 0, 255): 8,
    (0, 128, 128): 9,
}


def mask_to_labels(mask):
    m = np.asarray(mask)
    if m.ndim == 3:
        rgb = m[..., :3]
        labels = np.full(m.shape[:2], UNKNOWN, dtype=np.uint8)
        for color, label in RGB_PALETTE.items():
            labels[np.all(rgb == color, axis=2)] = label
        labels[np.all(rgb == (255, 255, 255), axis=2)] = IGNORE
    else:
        labels = m.astype(np.uint8)
    valid = (labels == IGNORE) | np.isin(labels, np.arange(len(CLASSES)))
    if not valid.all():
        raise ValueError(f"mask contains unrecognized/invalid values {np.unique(labels[~valid]).tolist()}")
    return labels

	"""Shared lossless mask taxonomy for the painting vision dataset."""
	import numpy as np

	CLASSES = ("other", "wall_unpainted", "wall_painted", "wall_uncertain", "skirting",
	"switch_outlet", "ac_unit", "door", "window", "wall_obstacle")
	IGNORE = 255
	UNKNOWN = 254

	# Visible surface material of wall pixels, predicted by the material head
	# independently of paint state. Wallpaper is the visible covering, so a
	# wallpapered gypsum wall is ``wallpaper``, not ``drywall``. Older datasets that
	# only use 0/1 stay valid; 255 marks unknown material or non-wall pixels.
	MATERIAL_CLASSES = ("other_wall_material", "drywall", "wallpaper")
	LEGACY_MATERIAL_CLASSES = MATERIAL_CLASSES[:2]
	# A smooth finished coat over board or plaster: the material is genuinely one of
	# the two and the photo cannot tell which. Trained toward 50/50 drywall/other, so
	# the model's honest answer there is low confidence (predict.py: "unknown").
	MATERIAL_DRYWALL_OR_OTHER = 254
	MATERIAL_VALUES = tuple(range(len(MATERIAL_CLASSES))) + (MATERIAL_DRYWALL_OR_OTHER, IGNORE)

	# Surface condition of decided wall pixels, from a third head: what has to be
	# prepared before paint goes on. ``paint_failure`` covers peeling, flaking,
	# blistering and cracked coatings; ``rust`` covers corrosion and rust staining.
	# Masks live in condition_masks/<split>/<stem>.png; 255 = unknown or not wall.
	# Datasets without condition masks are simply unsupervised for this head.
	CONDITION_CLASSES = ("sound", "paint_failure", "rust")
	CONDITION_VALUES = tuple(range(len(CONDITION_CLASSES))) + (IGNORE,)

	# Detection classes for the equipment around the wall: what has to be on site to
	# do the job, and what must not be painted over. These are objects, not pixel
	# labels - they never appear in the ten-class mask, so ``CLASSES`` above (and any
	# trained checkpoint) is unaffected by adding or reordering them. ``site_objects.py``
	# detects them without a model, from shape and colour cues.
	OBJECT_CLASSES = ("paint_can", "paint_tray", "brush", "roller", "ladder", "drop_sheet")

	# RGB values for lossless color-coded PNG masks.
	RGB_PALETTE = {
	(0, 0, 0): 0,
	(0, 0, 255): 1,
	(0, 255, 0): 2,
	(255, 255, 0): 3,
	(255, 0, 255): 4,
	(255, 0, 0): 5,
	(0, 255, 255): 6,
	(255, 128, 0): 7,
	(128, 0, 255): 8,
	(0, 128, 128): 9,
	}


	def mask_to_labels(mask):
	m = np.asarray(mask)
	if m.ndim == 3:
	rgb = m[..., :3]
	labels = np.full(m.shape[:2], UNKNOWN, dtype=np.uint8)
	for color, label in RGB_PALETTE.items():
	labels[np.all(rgb == color, axis=2)] = label
	labels[np.all(rgb == (255, 255, 255), axis=2)] = IGNORE
	else:
	labels = m.astype(np.uint8)
	valid = (labels == IGNORE) \| np.isin(labels, np.arange(len(CLASSES)))
	if not valid.all():
	raise ValueError(f"mask contains unrecognized/invalid values {np.unique(labels[~valid]).tolist()}")
	return labels