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_material_classification")

material_codes = {
    "AC": "Asphalt Concrete",
    "BL": "Block",
    "BR": "Brick",
    "CI": "Cast Iron",
    "CO": "Concrete",
    "CS": "Corrugated Steel",
    "DI": "Ductile Iron",
    "EP": "Epoxy",
    "GI": "Galvanized Iron",
    "MAR": "Masonry",
    "N": "Not specified, possibly a custom abbreviation",
    "OTH": "Other",
    "PE": "Polyethylene",
    "PF": "Plywood-Faced",
    "PP": "Polypropylene",
    "PVC": "Polyvinyl Chloride",
    "RC": "Reinforced Concrete",
    "ST": "Steel",
    "U": "Unspecified, possibly a custom abbreviation",
    "UPVC": "Unplasticized Polyvinyl Chloride",
    "VC": "Vinyl Coated",
    "XI": "Extra Impact",
    "XP": "Extruded Polystyrene",
    "Z": "Not specified, possibly a custom abbreviation"
}

material_full_names_list = [
    'Asphalt Concrete',
    'Block',
    'Brick',
    'Cast Iron',
    'Concrete',
    'Corrugated Steel',
    'Ductile Iron',
    'Epoxy',
    'Galvanized Iron',
    'Masonry',
    'Not specified, possibly a custom abbreviation',
    'Other',
    'Polyethylene',
    'Plywood-Faced',
    'Polypropylene',
    'Polyvinyl Chloride',
    'Reinforced Concrete',
    'Steel',
    'Unspecified, possibly a custom abbreviation',
    'Unplasticized Polyvinyl Chloride',
    'Vinyl Coated',
    'Extra Impact',
    'Extruded Polystyrene',
    'Not specified, possibly a custom abbreviation',
    'Vitrified Clay Lined'
]

def replace_label_with_full_name(res, defect_dict_key_code):
    new_res = {}
    for dic in res:
        # Splitting the label to handle possible suffix
        parts = dic["label"].split('_', 1)
        code = parts[0]
        suffix = '_' + parts[1] if len(parts) > 1 else ''

        # Replacing the code with its full name, if it exists in the dictionary
        full_name = defect_dict_key_code.get(code, code)

        # Constructing the new label with the suffix if it exists
        new_label = full_name + suffix
        new_res[new_label] = dic["score"]
    
    return new_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)
    
    # Extract labels and scores
    return replace_label_with_full_name(res, material_codes)


# 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="Material Classification with AI by Subterra",
    allow_flagging="manual",
    flagging_options=material_full_names_list
)

# Launch the interface
iface.launch()