from turtle import title
import gradio as gr
from transformers import pipeline
from PIL import Image

# Initialize the pipeline with your model
pipe = pipeline("image-classification", model="SubterraAI/ofwat_defects_classification")

defect_dict = {
    "CL": "Crack Longitudinal",
    "CLJ": "Crack Longitudinal at Joint",
    "CC": "Crack Circumferential",
    "CCJ": "Crack Circumferential at Joint",
    "CM": "Crack Multiple",
    "CMJ": "Crack Multiple at Joint",
    "CS": "Crack Spiral",
    "FL": "Fracture Longitudinal",
    "FC": "Fracture Circumferential",
    "FM": "Fracture Multiple",
    "FS": "Fracture Spiral",
    "B": "Broken",
    "BJ": "Broken Pipe at Joint",
    "H": "Hole",
    "D": "Deformation (Not Brick)",
    "XB": "Collapse",
    "JD": "Joint Displaced",
    "JDM": "Joint Displaced Medium",
    "JDL": "Joint Displaced Large",
    "OJ": "Open Joint",
    "OJM": "Open Joint Medium",
    "OJL": "Open Joint Large",
    "SW": "Increased Roughness",
    "SAV": "Visible Aggregate",
    "SAP": "Aggregate Projecting",
    "SRV": "Visible Reinforcement",
    "SRP": "Reinforcement Projecting",
    "SRC": "Corroded Reinforcement",
    "SS": "Surface Damage Spalling",
    "SZ": "Other Damage",
    "LX": "Line Defect",
    "WXC": "Weld Failure Circumferential",
    "WXL": "Weld Failure Longitudinal",
    "WXS": "Weld Failure Spiral",
    "RXM": "Defective Repair, part of wall missing",
    "RX": "Defective Repair",
    "DB": "Displaced bricks",
    "MB": "Missing bricks",
    "DI": "Dropped Invert",
    "EL": "Encrustation/Scale Light",
    "ESL": "Encrustation/Scale Light",
    "EM": "Encrustation/Scale Medium",
    "ESM": "Encrustation/Scale Medium",
    "EH": "Encrustation/Scale Heavy",
    "ESH": "Encrustation/Scale Heavy",
    "DEG": "Debris Grease",
    "DES": "Debris Silt",
    "RF": "Roots Fine",
    "RM": "Roots Mass",
    "RT": "Roots Tap",
    "IS(J)": "Infiltration Light Seeping Joint",
    "ID(J)": "Infiltration Light Dripping Joint",
    "IR(J)": "Infiltration Moderate Running Joint",
    "IRG(J)": "Infiltration Severe Gusher Joint",
    "CNI": "Connection Intruding",
    "CX": "Connection Defective",
    "SR": "Sealing Ring Intruding",
    "SRB": "Sealing Ring Broken",
    "SO": "Other sealent intruding",
    "CU": "Camera Underwater"
}

def replace_label_with_full_name(res, defect_dict_key_code):
    # Check if the label in the result is in the dictionary
    # If it is, replace it with the full name
    return {defect_dict_key_code.get(dic["label"], dic["label"]): dic["score"] for dic in res}


def classify_image(image):
    # Convert the input image to PIL format
    PIL_image = Image.fromarray(image).convert('RGB')
    
    # Classify the image using the pipeline
    res = pipe(PIL_image)

    replaced_res = replace_label_with_full_name(res, defect_dict)

    
    # Extract labels and scores
    return {dic["label"]: dic["score"] for dic in replaced_res}

# Create the Gradio interface
iface = gr.Interface(
    classify_image, 
    "image", 
    "label", 
    examples=[
        ["examples/CS.jpg"], 
        ["examples/GI.jpg"],
        ["examples/PP.jpg"],
        ["examples/RC.jpg"]
    ],
    description="Upload an image to classify its material.",
    title="Defects Classification with AI by Subterra"
    #allow_flagging="manual",
    #flagging_options=["obstruction", "no_obstruction"]
)

# Launch the interface
iface.launch()