A quick diagnostic for the WIK wafer dataset

process_wafers.py

import pickle
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from tqdm import tqdm, trange
from scipy.ndimage import zoom, binary_fill_holes

from skimage.color import label2rgb
from skimage.filters import sobel
from skimage.measure import label
from skimage.segmentation import watershed

def dilate_img(img):
    edges = sobel(img)
    markers = np.zeros_like(img)
    foreground, background = 1, 2
    markers[img < 30.0] = background
    markers[img > 100.0] = foreground
    ws = watershed(edges, markers)
    seg1 = label(ws == foreground)
    img = label2rgb(seg1,  bg_label=0)
    img = binary_fill_holes(img[...,0])
    return img

def calculate_2dft(img):
    ft = np.fft.ifftshift(img)
    ft = np.fft.fft2(ft)
    return np.fft.fftshift(ft)

with open('val_data.pkl','rb') as f:
    test_df = pickle.load(f)
defect_types = ['Center','Donut','Edge-Loc','Edge-Ring','Loc','Near-full','Random','Scratch','None']
N_IMAGES = test_df.waferMap.index.size
N_FRAMES = 100
fps = 1
snapshots = [zoom(test_df.waferMap.iloc[i], 1) for i in trange(N_FRAMES, desc='zoom')]
dil_snapshots = [dilate_img(snapshots[i]) for i in trange(N_FRAMES, desc='dilation')]
fft_snapshots = [np.log(abs(calculate_2dft(snapshots[i]))) for i in trange(N_FRAMES, desc='fft')]
codes = [test_df.failureCode.iloc[i] for i in range(N_FRAMES)]

# First set up the figure, the axis, and the plot element we want to animate
fig, axs = plt.subplots(2,2, figsize=(5,5), tight_layout=True)
(ax, ax_dil, ax_fft, ax_hist)  = axs.flatten()
a = snapshots[0]
fft_a = fft_snapshots[0]
dil_a = dil_snapshots[0]

im = ax.imshow(a, cmap='gray')
ax.set_title('Wafer Image')

im_dil = ax_dil.imshow(dil_a)
ax_dil.set_title('Dilation + Blobs')

im_fft = ax_fft.imshow(fft_a, cmap='gray')
ax_fft.set_title('Wafer FFT')

ax_hist.hist(a, bins='auto', fc='k', ec='k', lw=1, alpha=0.5)

pbar = tqdm(total=N_FRAMES, desc='render')
def animate_func(i):

    im.set_array(snapshots[i])
    im_fft.set_array(fft_snapshots[i])
    im_dil.set_array(dil_snapshots[i])

    
    ax_hist.clear()
    ax_hist.hist(snapshots[i],bins='auto', fc='k', ec='k', lw=1, alpha=0.5)
    ax_hist.set_ylim([0,100])
    ax_hist.set_title('Histogram')

    im.get_figure().suptitle(f'Index {i} - FailureType {defect_types[codes[i]]}')
    pbar.update(1)
    return [im]

anim = animation.FuncAnimation(
    fig, 
    animate_func, 
    frames = N_FRAMES,
    interval = 1000 / fps, # in ms
)
# plt.show()
anim.save('wafers.mp4', fps=fps, extra_args=['-vcodec', 'libx264'], dpi=300)
pbar.close()