hsi_painter.py

import numpy as np
import spectral.io.envi as envi
from matplotlib import pyplot as plt

"""
超高速小畫家

* 顏色增加改for迴圈寫法

* 改顏色也換寫法

* 消除會越畫越慢的問題: draw() -> draw_artist() and blit()
ref:https://stackoverflow.com/questions/8955869/why-is-plotting-with-matplotlib-so-slow

* 調整間距的快速方法 ax.relim() and ax.autoscale_view()
ref:https://stackoverflow.com/questions/31322394/dynamically-change-y-limit
"""

# Global variable
current_color_idx = 0  # 建立全域的筆刷顏色
current_brush = None  # 建立全域的筆刷的class紀錄
brush_list = []  # 建立所有使用的筆刷


class BrushBuilder:
    lock = None  # 建立一個全域的鎖，當點擊到畫布上就鎖住，放開滑鼠就變回None
    first_draw = True

    def __init__(self, fig, axes_canvas, line_canvas):
        self.fig = fig
        self.axes_canvas = axes_canvas  # 左邊畫點畫布
        self.line_canvas = line_canvas  # 右邊畫線畫布

        # 建立陣列紀錄x,y的座標
        self.xs = list(self.axes_canvas.get_xdata())
        self.ys = list(self.axes_canvas.get_ydata())

        # 綁定事件
        self.connect()

    def get_color(self):  # 拿當下畫筆顏色(key_press改顏色需要)
        assert self.axes_canvas.get_color() == self.line_canvas.get_color()
        return self.axes_canvas.get_color()

    def connect(self):  # 綁定滑鼠事件： 點擊、移動、放開
        self.cidp = self.fig.canvas.mpl_connect("button_press_event", self.on_press)
        self.cidm = self.fig.canvas.mpl_connect("motion_notify_event", self.on_motion)
        self.cidr = self.fig.canvas.mpl_connect("button_release_event", self.on_release)

    def disconnect(self):  # 解除綁定
        self.fig.canvas.mpl_disconnect(self.cidp)
        self.fig.canvas.mpl_disconnect(self.cidm)
        self.fig.canvas.mpl_disconnect(self.cidr)

    def clear(self):  # 清空座標與光譜並重畫
        self.xs.clear()
        self.ys.clear()
        self.axes_canvas.set_data([], [])
        self.line_canvas.set_data([], [])
        self.fig.canvas.draw()

    def get_xy(self):
        return (self.ys, self.xs)

    def on_press(self, event):  # 滑鼠點擊事件
        if event.inaxes != self.axes_canvas.axes:
            return  # 如果滑鼠不在畫布內就返回
        if BrushBuilder.lock is not None:
            return  # 如果沒有放開過滑鼠就返回
        BrushBuilder.lock = self  # 點擊畫布，鎖上
        self.on_draw(event.xdata, event.ydata)  # 畫圖

    def on_motion(self, event):  # 滑鼠移動事件
        if event.inaxes != self.axes_canvas.axes:
            return  # 如果滑鼠不在畫布內就返回
        if BrushBuilder.lock is not self:
            return  # 如果沒有按著畫布就返回
        self.on_draw(event.xdata, event.ydata)  # 畫圖

    def on_release(self, event):  # 滑鼠放開事件
        if BrushBuilder.lock is not self:
            return  # 如果沒點過畫布就返回
        BrushBuilder.lock = None  # 滑鼠放開，把鎖釋放

    def on_draw(self, x, y):  # 畫點事件
        # 紀錄x,y資訊
        x = int(x)
        y = int(y)
        self.xs.append(x)
        self.ys.append(y)

        # 畫光譜
        tmp = mat[y, x]
        # tmp = np.gradient(tmp)
        self.line_canvas.set_data(range(len(tmp)), tmp)
        self.line_canvas.set_color(self.get_color())
        # self.line_canvas.axes.relim()
        # self.line_canvas.axes.autoscale()
        self.line_canvas.axes.draw_artist(self.line_canvas)

        # 畫座標
        self.axes_canvas.set_data(self.xs, self.ys)
        self.axes_canvas.axes.draw_artist(self.axes_canvas)

        # 更新畫布
        self.fig.canvas.blit()

        if self.first_draw:
            self.fig.canvas.draw()
            self.first_draw = False


def cem(HIM, d):
    w, h, n = HIM.shape
    r = HIM.reshape(-1, n).T
    R = 1 / (w * h) * r @ r.T
    Rinv = np.linalg.inv(R)
    res = (r.T @ Rinv @ d.T) / (d @ Rinv @ d.T)
    return res.reshape(w, h)


# def glcm(HIM, )


def key_press(event):
    global current_brush, current_color_idx, mat, ax_axes, name  # 使用到全域，先宣告

    if event.key == "x":  # 'x' 切換顏色
        background = current_brush.fig.canvas.copy_from_bbox(
            current_brush.line_canvas.axes.bbox
        )  # 保留畫線的數據
        current_brush.disconnect()  # 先解除綁定

        # 改變顏色
        current_color_idx = current_color_idx + 1
        if current_color_idx == len(brush_list):
            current_color_idx = 0
        current_brush = brush_list[current_color_idx]

        current_brush.connect()  # 重新綁定
        current_brush.axes_canvas.axes.set_title(
            "Press 'x' change color (Current: %s)" % (current_brush.get_color())
        )  # 換標題
        current_brush.fig.canvas.draw()  # 更新畫布
        current_brush.fig.canvas.restore_region(background)  # 原本的資料保留
    if event.key == "c":  # 'c': 清空畫布
        for b in brush_list:  # 清空全部畫筆內的座標與光譜
            b.clear()
    if event.key == "e":  # 'e': 跑CEM
        print("run cem")
        xy = brush_list[0].get_xy()
        res = cem(mat, mat[xy].mean(0))

        res = np.abs(res)

        # th = get_thres_n_times(res, 1)
        # res[res < th] = 0
        res[res >= 1] = 1
        fig = plt.figure()
        ax = fig.add_subplot(111)
        im = ax.imshow(res)
        fig.colorbar(im)
        ax.autoscale(True)
        ax.axis("off")
        fig.savefig(f"{name}_cem", bbox_inches="tight", dpi=200)
        fig.show()

        # ax_axes.figure.savefig(f"{name}_img", bbox_inches="tight", dpi=300)
        extent = ax_axes.get_window_extent().transformed(ax_axes.figure.dpi_scale_trans.inverted())
        ax_axes.autoscale(True)
        ax_axes.axis("off")
        ax_axes.figure.savefig(f"{name}_img", bbox_inches=extent, dpi=150)


def get_thres_n_times(img, n):
    from skimage.filters import threshold_otsu

    img = img.reshape(-1)
    for i in range(n - 1):
        thresh = threshold_otsu(img)
        bimage = img.copy()

        bimage[bimage < thresh] = 0
        bimage[bimage >= thresh] = 1

        zero_index = np.argwhere(bimage == 0)
        zero_index = list(zero_index[:, 0])

        img = np.delete(img, zero_index, 0)

    threshold = threshold_otsu(img)

    return threshold


def load_envi(path):
    mat = envi.open(path + ".hdr", path + ".raw")
    try:
        arr = mat.asarray()
    except:
        arr = mat.open_memmap(interleave="source", writable=True)
        arr = mat.asarray()
    return mat.metadata["wavelength"], arr


###########顯示的圖，改這邊就好###########
name = "核桃_RT"
wave, mat = load_envi(r"C:\Users\user\Desktop\新增資料夾\No-2_20-20-1_RT")

# 要畫資料改這邊
draw_which_band = 98  # 要顯示哪個Band
alpha = 0.3  # 線跟點的透明度

###########以下為組態區###########
brush_colors = ["red", "blue", "green", "purple", "cyan", "black", "magenta"]  # 顏色組
# 建立畫布
fig = plt.figure(figsize=[12, 6])
fig.canvas.mpl_connect("key_press_event", key_press)  # 綁定按鍵監聽給視窗

ax_axes = fig.add_subplot(121)  # 原圖畫布
ax_axes.imshow(mat[..., draw_which_band], "jet")  # 畫某個Band
ax_line = fig.add_subplot(122)  # 光譜畫布

# 建立筆刷並記錄
for color in brush_colors:
    (axes_canvas,) = ax_axes.plot(
        [], [], alpha=alpha, color=color, marker="o", linestyle="none"
    )
    (line_canvas,) = ax_line.plot([], [], alpha=alpha, color=color)
    ax_line.set_ylim((0, 1))
    ax_line.set_xlim((-5, mat.shape[-1] + 10))

    brush_list.append(BrushBuilder(fig, axes_canvas, line_canvas))

current_brush = brush_list[current_color_idx]  # 預設為第一個顏色
ax_axes.set_title("Press 'x' change color (Current: %s)" % (current_brush.get_color()))
ax_line.set_title("Press 'c' clear plot")
ax_line.set_xticklabels(np.asfarray(wave).astype("i"))

plt.tight_layout()
plt.show()