Add various plotting utils for easy comparison

estimate_albedo.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python3
+import argparse
+from pathlib import Path
+
+import numpy as np
+from PIL import Image
+import matplotlib.pyplot as plt
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+
+def linearize_image(img_pil):
+    """Convert a PIL RGB image (sRGB) to linear RGB in [0,1]."""
+    img_srgb = np.array(img_pil.convert("RGB"), dtype=np.float32) / 255.0
+    mask = img_srgb <= 0.04045
+    img_lin = np.where(mask, img_srgb / 12.92, ((img_srgb + 0.055) / 1.055) ** 2.4)
+    return np.clip(img_lin, 0, 1)
+
+
+def luminance(img_lin):
+    """Compute human-eye weighted luminance (ITU-R BT.709) from linear RGB."""
+    Y = 0.2126 * img_lin[..., 0] + 0.7152 * img_lin[..., 1] + 0.0722 * img_lin[..., 2]
+    return np.clip(Y, 0, 1)
+
+
+def lightness_map(img_pil):
+    """Compute simple lightness (HSL L) from a PIL RGB image."""
+    return linearize_image(img_pil).mean(axis=-1)
+
+
+def value_map(img_pil):
+    """Compute simple brightness (HSV V) from a PIL RGB image."""
+    return linearize_image(img_pil).max(axis=-1)
+
+
+def luminance_map(img_pil):
+    """Compute human-eye weighted luminance (ITU-R BT.709) from a PIL RGB image."""
+    return luminance(linearize_image(img_pil))
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Run Mask2Former segmentation and save material overlay."
+    )
+    parser.add_argument(
+        "image_path",
+        type=Path,
+        help="Path to input image",
+    )
+    args = parser.parse_args()
+
+    image_path: Path = args.image_path
+    return image_path
+
+
+def main():
+    # Parse args
+    image_path = parse_args()
+
+    # Open image
+    image = Image.open(image_path).convert("RGB")
+
+    # Calculate albedo
+    luminance = luminance_map(image)
+    albedo = luminance / luminance.max()
+
+    # Create image
+    fig, ax = plt.subplots(figsize=(12, 12))
+    im = ax.imshow(albedo, vmin=0, vmax=1, cmap="pink")
+    ax.set_title("Luminace as proxy for albedo")
+    ax.axis("off")
+
+    # Attach colorbar that matches height of axes
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes("right", size="5%", pad=0.05)  # adjust size and pad
+    plt.colorbar(im, cax=cax)
+
+    output_path = Path.cwd() / (image_path.stem + "_albedo.png")
+    fig.savefig(output_path, bbox_inches="tight")
+
+
+if __name__ == "__main__":
+    plt.rcParams.update({"font.size": 24})
+    main()

identify_materials.py

@@ -64,7 +64,7 @@
 
 # Label map
 id2label = {
-    0: "Background",
+    0: "Background/Unclassified",
     1: "Wood/Bamboo",
     2: "Ground tile",
     3: "Brick",
@@ -110,11 +110,81 @@
 }
 
 
+def plot_overlay(image, mask):
+    fig, ax = plt.subplots(figsize=(12, 12))
+
+    # Map semantic mask to RGB
+    h, w = mask.shape
+    color_mask = np.zeros((h, w, 3), dtype=np.uint8)
+    for lbl, color in enumerate(label_colors):
+        color_mask[mask == lbl] = color
+
+    # Overlay
+    image_np = np.array(image.resize((w, h)))
+    overlay = (0.5 * image_np + 0.5 * color_mask).astype(np.uint8)
+
+    # Plot
+    ax.imshow(overlay)
+    ax.axis("off")
+
+    # Legend
+    classes = np.insert(np.unique(mask), 0, 0)
+    legend_handles = [
+        Patch(color=np.array(c) / 255.0, label=id2label[i])
+        for i, c in enumerate(label_colors)
+        if i in classes
+    ]
+    plt.legend(
+        handles=legend_handles,
+        bbox_to_anchor=(1, 1),
+        loc="upper left",
+        fontsize="small",
+        frameon=False,
+    )
+    ax.set_title("Material estimates by BFMS")
+    return fig, ax
+
+
+def plot_side_by_side(image, mask):
+    # Map semantic mask to RGB
+    h, w = mask.shape
+    color_mask = np.zeros((h, w, 3), dtype=np.uint8)
+    for lbl, color in enumerate(label_colors):
+        color_mask[mask == lbl] = color
+
+    fig, axs = plt.subplots(1, 2, figsize=(15, 10))
+    axs[0].imshow(image)
+    axs[1].imshow(color_mask)
+
+    # Legend
+    classes = np.insert(np.unique(mask), 0, 0)
+    legend_handles = [
+        Patch(color=np.array(c) / 255.0, label=id2label[i])
+        for i, c in enumerate(label_colors)
+        if i in classes
+    ]
+    plt.legend(
+        handles=legend_handles,
+        bbox_to_anchor=(1.05, 1),
+        loc="upper left",
+        fontsize="small",
+    )
+
+    for ax in axs:
+        ax.axis("off")
+    fig.tight_layout()
+    return fig, axs
+
+
 def main():
     parser = argparse.ArgumentParser(
         description="Run Mask2Former segmentation and save material overlay."
     )
-    parser.add_argument("image_path", type=Path, help="Path to input image")
+    parser.add_argument(
+        "image_path",
+        type=Path,
+        help="Path to input image",
+    )
     parser.add_argument(
         "--model_path",
         type=Path,
@@ -160,42 +230,21 @@ def main():
     pixel_class_probs = torch.einsum("qc,qhw->chw", class_probs, masks_probs)
     semantic_mask = torch.argmax(pixel_class_probs, dim=0).cpu().numpy()
 
-    # Map semantic mask to RGB
-    h, w = semantic_mask.shape
-    color_mask = np.zeros((h, w, 3), dtype=np.uint8)
-    for lbl, color in enumerate(label_colors):
-        color_mask[semantic_mask == lbl] = color
-
-    # Overlay
-    image_np = np.array(image.resize((w, h)))
-    overlay = (0.5 * image_np + 0.5 * color_mask).astype(np.uint8)
-
-    # Plot
-    plt.figure(figsize=(12, 12))
-    plt.imshow(overlay)
-    plt.axis("off")
-
-    # Legend
-    classes = np.unique(semantic_mask)
-    legend_handles = [
-        Patch(color=np.array(c) / 255.0, label=id2label[i])
-        for i, c in enumerate(label_colors)
-        if i in classes
-    ]
-    plt.legend(
-        handles=legend_handles,
-        bbox_to_anchor=(1.05, 1),
-        loc="upper left",
-        fontsize="small",
-    )
-    plt.tight_layout()
-
-    # Output filename
-    output_filename = f"{image_path.stem}_materials.png"
+    # Overlay image
+    fig, ax = plot_overlay(image, semantic_mask)
+    output_filename = f"{image_path.stem}_materials_overlay.png"
     output_path = Path.cwd() / output_filename
-    plt.savefig(output_path)
+    fig.savefig(output_path, bbox_inches="tight")
     print(f"Saved overlay to {output_path}")
 
+    # Side by side image
+    fig, axs = plot_side_by_side(image, semantic_mask)
+    output_filename = f"{image_path.stem}_materials_sidebyside.png"
+    output_path = Path.cwd() / output_filename
+    fig.savefig(output_path, bbox_inches="tight")
+    print(f"Saved side-by-side image to {output_path}")
+
 
 if __name__ == "__main__":
+    plt.rcParams.update({"font.size": 24})
     main()

plot_original.py

@@ -0,0 +1,41 @@
+import argparse
+from pathlib import Path
+from PIL import Image
+from matplotlib import pyplot as plt
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Run Mask2Former segmentation and save material overlay."
+    )
+    parser.add_argument(
+        "image_path",
+        type=Path,
+        help="Path to input image",
+    )
+    args = parser.parse_args()
+
+    image_path: Path = args.image_path
+    return image_path
+
+
+def main():
+    # Parse args
+    image_path = parse_args()
+
+    # Open image
+    image = Image.open(image_path).convert("RGB")
+
+    # Create image
+    fig, ax = plt.subplots(figsize=(12, 12))
+    ax.imshow(image)
+    ax.set_title("Original image")
+    ax.axis("off")
+
+    output_path = Path.cwd() / (image_path.stem + "_original.png")
+    fig.savefig(output_path, bbox_inches="tight")
+
+
+if __name__ == "__main__":
+    plt.rcParams.update({"font.size": 24})
+    main()

pyproject.toml

@@ -5,6 +5,7 @@ description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.11"
 dependencies = [
+    "ipython>=9.4.0",
     "matplotlib>=3.10.5",
     "safetensors>=0.6.2",
     "scipy>=1.16.1",