Added the full name in the prediction

app.py

@@ -6,6 +6,74 @@ 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
@@ -13,9 +81,12 @@ def classify_image(image):
     
     # 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 res}
+    return {dic["label"]: dic["score"] for dic in replaced_res}
 
 # Create the Gradio interface
 iface = gr.Interface(
@@ -30,6 +101,8 @@ iface = gr.Interface(
     ],
     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