Wafer bin map

Python Code

"""Semiconductor — Wafer bin map (WaferGuard dataset)."""
import io
import tempfile


import numpy as np
import plotly.graph_objects as go
import requests
import warnings

warnings.filterwarnings("ignore")
requests.packages.urllib3.disable_warnings()  # type: ignore


URL = "https://huggingface.co/datasets/oliversinn/waferguard-dataset/resolve/main/data/train-00000-of-00002.parquet"


DEFECT_NAMES = [
    "Center", "Donut", "Edge-Loc", "Edge-Ring",
    "Loc", "Near-full", "Random", "Scratch",
]


def load_wafer():
    """Try to load real wafer data, fall back to synthetic."""
    try:
        import pyarrow.parquet as pq

        print("  downloading parquet ...")
        r = requests.get(URL, verify=False, timeout=60)
        r.raise_for_status()
        with tempfile.NamedTemporaryFile(suffix=".parquet", delete=False) as f:
            f.write(r.content)
            tmp = f.name

        table = pq.read_table(tmp)
        df = table.to_pandas()
        print(f"  columns: {list(df.columns)}")
        print(f"  shape: {df.shape}")

        # image: each row is a (52,) object array of (52,) int64 arrays
        # label: each row is an 8-element multi-hot numpy array
        # Find first row with a defect die (value == 2 in image)
        found_idx = None
        for i in range(min(500, len(df))):
            try:
                img_row = df["image"].iloc[i]
                arr_test = np.array([list(r) for r in img_row], dtype=float)
                if 2.0 in arr_test:
                    found_idx = i
                    break
            except Exception:
                pass
        if found_idx is None:
            found_idx = 0
        print(f"  using row {found_idx}")

        row = df.iloc[found_idx]
        img = row["image"]

        # Build 2D array from list-of-arrays structure
        arr = np.array([list(r) for r in img], dtype=float)
        print(f"  image shape: {arr.shape}, unique: {np.unique(arr)}")

        # Decode label: multi-hot 8-bit
        lbl_arr = row["label"]
        active = [DEFECT_NAMES[i] for i, v in enumerate(lbl_arr) if v]
        label = "+".join(active) if active else "Normal"
        print(f"  label: {label}")

        return arr, label

    except Exception as e:
        print(f"  real data failed ({e}), using synthetic wafer")
        return None, None


def synthetic_wafer(size=52):
    """Generate a synthetic wafer with a ring-defect pattern."""
    arr = np.zeros((size, size), dtype=float)
    cx, cy = size / 2, size / 2
    r_wafer = size / 2 - 1

    for i in range(size):
        for j in range(size):
            r = np.sqrt((i - cx) ** 2 + (j - cy) ** 2)
            if r > r_wafer:
                arr[i, j] = np.nan          # off-wafer
            else:
                arr[i, j] = 1               # pass (green)

    # Add a ring-shaped defect cluster
    np.random.seed(42)
    for i in range(size):
        for j in range(size):
            r = np.sqrt((i - cx) ** 2 + (j - cy) ** 2)
            # Ring defect at r ≈ 18-22
            if 18 < r < 22 and not np.isnan(arr[i, j]):
                if np.random.random() < 0.7:
                    arr[i, j] = 2           # fail (red)

    # Add some center dies
    for i in range(size):
        for j in range(size):
            r = np.sqrt((i - cx) ** 2 + (j - cy) ** 2)
            if r < 5 and not np.isnan(arr[i, j]):
                if np.random.random() < 0.4:
                    arr[i, j] = 2

    return arr, "Edge Ring"


def apply_circular_mask(arr):
    """Set off-wafer cells (value == 0 at boundary) to NaN if not already done."""
    size = arr.shape[0]
    cx, cy = size / 2, size / 2
    r_wafer = size / 2 - 1
    masked = arr.copy().astype(float)
    for i in range(size):
        for j in range(size):
            r = np.sqrt((i - cx) ** 2 + (j - cy) ** 2)
            if r > r_wafer:
                masked[i, j] = np.nan
    return masked


def generate():
    arr, label = load_wafer()

    if arr is None:
        arr, label = synthetic_wafer()
        used_synthetic = True
    else:
        # Apply circular mask: 0 = off-wafer → NaN
        arr = apply_circular_mask(arr)
        used_synthetic = False

    # Replace 0 (off-wafer) with NaN
    arr[arr == 0] = np.nan

    size = arr.shape[0]
    print(f"  wafer size: {size}x{size}, label: {label}")
    print(f"  unique values (excl NaN): {np.unique(arr[~np.isnan(arr)])}")

    # Discrete colorscale: 1 = pass (green), 2 = fail (red)
    colorscale = [
        [0.0, "#55B685"],   # pass — GREEN
        [0.5, "#55B685"],
        [0.5, "#DA5597"],   # fail — PINK/red
        [1.0, "#DA5597"],
    ]

    fig = go.Figure(go.Heatmap(
        z=arr,
        colorscale=colorscale,
        zmin=1,
        zmax=2,
        colorbar=dict(
            tickvals=[1, 2],
            ticktext=["Pass", "Fail"],
            title="Die status",
            thickness=16,
        ),
        showscale=True,
        hovertemplate="Row: %{y}<br>Col: %{x}<br>Status: %{z}<extra></extra>",
        xgap=1,
        ygap=1,
    ))

    fig.update_layout(
        xaxis=dict(
            showticklabels=False,
            showgrid=False,
            zeroline=False,
            constrain="domain",
        ),
        yaxis=dict(
            showticklabels=False,
            showgrid=False,
            zeroline=False,
            scaleanchor="x",
            constrain="domain",
            autorange="reversed",
        ),
        margin=dict(t=40, b=40, l=40, r=80),
        annotations=[
            dict(
                text=f"Defect class: {label}" + (" (synthetic)" if used_synthetic else ""),
                x=0.5, y=1.04,
                xref="paper", yref="paper",
                showarrow=False,
                font=dict(size=11, color=MUTED),
            )
        ],
    )

    return fig


if __name__ == "__main__":
    generate()

Made with Plotly