app.py

import numpy as np
import pandas as pd
import gradio as gr
import cv2
from tensorflow import keras as k

# ============================== variables ==============================
image_size = 256
num_classes = 3

in_channel_tool = 3
in_channel_spec = 9
in_channel_scal = 9
in_channel_chip = 3
in_channel_work = 3
img_rows, img_cols = image_size, image_size

# ============================== paths ==============================
model_class_path = f"Models/siren3d_v3_hexa_base_multi_tsscw.h5"
model_reg_path = f"Models/siren3d_v3_hexa_regression_multi_tsscw.h5"

# dataset
csv_path = "Dataset/labels_sample.csv"
csv_path_reg = "Dataset/labels_reg_sample.csv"
tool_path = "Dataset/tool"
spec_path = "Dataset/spec"
scal_path = "Dataset/scal"
chip_path = "Dataset/chip"
work_path = "Dataset/work"

# ============================== load/prepare dataset ==============================
df = pd.read_csv(csv_path)
df_reg = pd.read_csv(csv_path_reg)

df["tool"] = df.id.map(lambda id: f"{tool_path}/{id}.jpg")
df["spec_x"] = df.id.map(lambda id: f"{spec_path}/x/{id}.jpg")
df["spec_y"] = df.id.map(lambda id: f"{spec_path}/y/{id}.jpg")
df["spec_z"] = df.id.map(lambda id: f"{spec_path}/z/{id}.jpg")
df["scal_x"] = df.id.map(lambda id: f"{scal_path}/x/{id}.png")
df["scal_y"] = df.id.map(lambda id: f"{scal_path}/y/{id}.png")
df["scal_z"] = df.id.map(lambda id: f"{scal_path}/z/{id}.png")
df["chip"] = df.id.map(lambda id: f"{chip_path}/{id}.jpg")
df["work"] = df.id.map(lambda id: f"{work_path}/{id}.png")

# ============================== examples ===========================================
exs = []
for i in range(len(df)):
    row = df.iloc[i,:]
    tool_id = row.id
    image_label = row.image_label 
    tool = row.tool 
    spec_x = row.spec_x 
    spec_y = row.spec_y 
    spec_z = row.spec_z 
    scal_x = row.scal_x 
    scal_y = row.scal_y 
    scal_z = row.scal_z 
    chip = row.chip 
    work = row.work 

    task = "Regression"
    if i % 2 == 0:
        task = "Classification"
    example = [tool_id, image_label, task, tool, spec_x, spec_y, spec_z, scal_x, scal_y, scal_z, chip, work]
    exs.append(example)

# ============================== preprocess ===========================================
def process_img(img, img_rows, img_cols, channels):
    """
    Reads the spectogram files from disk and normalizes the pixel values
      @params:
        img - Data of the image
        img_rows - The image height.
        img_cols - The image width.
        as_grey - Read the image as Greyscale or RGB.
        channels - Number of channels.
      @returns:
        The created and compiled model (Model)
    """
    img = cv2.imread(img)
    img = cv2.resize(img, dsize=(img_rows, img_cols), interpolation=cv2.INTER_CUBIC)
    img = np.asarray(img, dtype=np.float32)

    # normalize
    #print(np.max(img))
    img = img / 255.0 #np.max(img)

    # reshape to match Keras expectaions
    img = img.reshape(img_rows, img_cols, channels)

    return img

def process_specs(img_x, img_y, img_z, img_rows, img_cols, channels):
    img_x = cv2.imread(img_x)
    img_y = cv2.imread(img_y)
    img_z = cv2.imread(img_z)

    img = []
    img_x = cv2.resize(img_x, dsize=(img_rows, img_cols), interpolation=cv2.INTER_CUBIC)
    img_y = cv2.resize(img_y, dsize=(img_rows, img_cols), interpolation=cv2.INTER_CUBIC)
    img_z = cv2.resize(img_z, dsize=(img_rows, img_cols), interpolation=cv2.INTER_CUBIC)
    img.append([img_x, img_y, img_z])
    img = np.asarray(img, dtype=np.float32)

    # normalize
    #print(np.max(img))
    img = img / 255.0 #np.max(img)

    # reshape to match Keras expectaions
    img = img.reshape(img_rows, img_cols, channels)

    return img


# ============================== load model ==============================================
model_class = k.models.load_model(model_class_path, compile=False)
model_reg = k.models.load_model(model_reg_path, compile=False)
    
# ============================== app =====================================================

def change_output_labels(choice):
  if choice == "Classification":
    return [
        gr.Label(value=None, label="Actual Label", visible=True), gr.Label(value=None, label="Predicted Label", visible=True), 
        gr.Label(label="Actual Gaps", visible=False), gr.Label(label="Predicted Gaps", visible=False),
        gr.Label(label="Actual Flank Wear", visible=False), gr.Label(label="Predicted Flank Wear", visible=False),
        gr.Label(label="Actual Overhang", visible=False), gr.Label(label="Predicted Overhang", visible=False)
      ]
  else:
    return [
        gr.Label("Sharp", label="Actual Label", visible=False), gr.Label("Sharp", label="Predicted Label", visible=False), 
        gr.Label(value=None, label="Actual Gaps", visible=True), gr.Label(value=None, label="Predicted Gaps", visible=True),
        gr.Label(value=None, label="Actual Flank Wear", visible=True), gr.Label(value=None, label="Predicted Flank Wear", visible=True),
        gr.Label(value=None, label="Actual Overhang", visible=True), gr.Label(value=None, label="Predicted Overhang", visible=True)
      ]


def predict(tool_id, task, label, tool, spec_x, spec_y, spec_z, scal_x, scal_y, scal_z, chip, work):
    if task is None:
      raise gr.Error("Choose a task first!")

    labels = ['sharp', 'used', 'dulled']
    tool = process_img(tool, img_rows, img_cols, in_channel_tool)
    spec = process_specs(spec_x, spec_y, spec_z, img_rows, img_cols, in_channel_spec)
    scal = process_specs(scal_x, scal_y, scal_z, img_rows, img_cols, in_channel_scal)
    chip = process_img(chip, img_rows, img_cols, in_channel_chip)
    work = process_img(work, img_rows, img_cols, in_channel_work)


    inputs = [np.array([tool,]), np.array([spec,]), np.array([scal,]), np.array([chip,]), np.array([work,])]
    print(task)
    if task == "Classification":
      y_score = model_class.predict(inputs)
      y_pred = {label:float(score) for label, score in zip(labels, y_score[0])}
      return [
          gr.Label(value=label, label="Actual Label", visible=True), gr.Label(value=y_pred, label="Predicted Label", visible=True), 
          gr.Label(label="Actual Gaps", visible=False), gr.Label(label="Predicted Gaps", visible=False),
          gr.Label(label="Actual Flank Wear", visible=False), gr.Label(label="Predicted Flank Wear", visible=False),
          gr.Label(label="Actual Overhang", visible=False), gr.Label(label="Predicted Overhang", visible=False)
        ]
    else:
      y_score = model_reg.predict(inputs)
      print(y_score)
      gaps_pred = str(y_score[0][0])
      flank_wear_pred = str(y_score[0][1])
      overhang_pred = str(y_score[0][2])

      actual = df_reg[df_reg["id"] == tool_id].values
      gaps_actual = str(actual[0][1])
      flank_wear_actual = str(actual[0][2])
      overhang_actual = str(actual[0][3])
      print(gaps_actual)
      return [
        gr.Label("Sharp", label="Actual Label", visible=False), gr.Label("Sharp", label="Predicted Label", visible=False), 
        gr.Label(value=gaps_actual, label="Actual Gaps", visible=True), gr.Label(value=gaps_pred, label="Predicted Gaps", visible=True),
        gr.Label(value=flank_wear_actual, label="Actual Flank Wear", visible=True), gr.Label(value=flank_wear_pred, label="Predicted Flank Wear", visible=True),
        gr.Label(value=overhang_actual, label="Actual Overhang", visible=True), gr.Label(value=overhang_pred, label="Predicted Overhang", visible=True)
      ]

# ============================== blocks ========================================================
title = r"""
<h1 align="center">Impala</h1>
"""
description = r"""
<b>Official 🤗 Gradio demo</b> for <a href='https://github.com/hubtru/Impala' 
target='_blank'><b>Expandable Isotropic Multimodal Patch Learning Neural Architecture for the Hexa-modal (9) time-series and images data</b></a>.<br>
"""

with gr.Blocks() as demo:
  gr.Markdown(value=title)
  gr.Markdown(description)
  with gr.Row():
    with gr.Column():
      with gr.Row():
        tool_id = gr.Textbox("T1R2B1", label="Tool")
        label_input = gr.Textbox("Sharp", label="Label")
      task = gr.Radio(["Classification", "Regression"], label="Task")
      with gr.Row():
        tool = gr.Image(label="Tool", type="filepath")
      with gr.Row():
        chip = gr.Image(label="Chip", type="filepath")
        work = gr.Image(label="Work", type="filepath")
      with gr.Row():
        spec_x = gr.Image(label="Spec_x", type="filepath")
        spec_y = gr.Image(label="Spec_y", type="filepath")
        spec_z = gr.Image(label="Spec_z", type="filepath")
      with gr.Row():
        scal_x = gr.Image(label="Scal_x", type="filepath")
        scal_y = gr.Image(label="Scal_y", type="filepath")
        scal_z = gr.Image(label="Scal_z", type="filepath")
      submit_btn = gr.Button("Submit", variant="primary")
      
    with gr.Column():
      output_labels = [
        gr.Label("Sharp", label="Actual Label"),
        gr.Label("Sharp", label="Predicted Label"),
        gr.Label(label="Actual Gaps", visible=False), 
        gr.Label(label="Predicted Gaps", visible=False), 
        gr.Label(label="Actual Flank Wear", visible=False), 
        gr.Label(label="Predicted Flank Wear", visible=False), 
        gr.Label(label="Actual Overhang", visible=False), 
        gr.Label(label="Predicted Overhang", visible=False), 
      ]
  
  examples = gr.Examples(examples=exs, inputs=[tool_id, label_input, task, tool, spec_x, spec_y, spec_z, scal_x, scal_y, scal_z, chip, work])
  task.change(fn=change_output_labels, inputs=task, outputs=output_labels)
  submit_btn.click(fn=predict, inputs=[tool_id, task, label_input, tool, spec_x, spec_y, spec_z, scal_x, scal_y, scal_z, chip, work], outputs=output_labels)
demo.launch()