|
import gradio as gr |
|
from transformers import pipeline |
|
from PIL import Image |
|
import os |
|
from huggingface_hub import HfApi, upload_file |
|
import io |
|
import numpy as np |
|
import uuid |
|
|
|
|
|
pipe = pipeline("image-classification", model="SubterraAI/ofwat_defects_classification") |
|
HF_TOKEN = os.getenv('HF_TOKEN') |
|
DATASET_NAME = "SubterraAI/ofwat_defects_loop" |
|
hf_api = HfApi() |
|
|
|
|
|
flagged_data_dir = "./flagged_data" |
|
|
|
defect_full_names_list = [ |
|
"Crack Longitudinal", |
|
"Crack Longitudinal at Joint", |
|
"Crack Circumferential", |
|
"Crack Circumferential at Joint", |
|
"Crack Multiple", |
|
"Crack Multiple at Joint", |
|
"Crack Spiral", |
|
"Fracture Longitudinal", |
|
"Fracture Circumferential", |
|
"Fracture Multiple", |
|
"Fracture Spiral", |
|
"Broken", |
|
"Broken Pipe at Joint", |
|
"Hole", |
|
"Deformation (Not Brick)", |
|
"Collapse", |
|
"Joint Displaced", |
|
"Joint Displaced Medium", |
|
"Joint Displaced Large", |
|
"Open Joint", |
|
"Open Joint Medium", |
|
"Open Joint Large", |
|
"Increased Roughness", |
|
"Visible Aggregate", |
|
"Aggregate Projecting", |
|
"Visible Reinforcement", |
|
"Reinforcement Projecting", |
|
"Corroded Reinforcement", |
|
"Surface Damage Spalling", |
|
"Other Damage", |
|
"Line Defect", |
|
"Weld Failure Circumferential", |
|
"Weld Failure Longitudinal", |
|
"Weld Failure Spiral", |
|
"Defective Repair, part of wall missing", |
|
"Defective Repair", |
|
"Displaced bricks", |
|
"Missing bricks", |
|
"Dropped Invert", |
|
"Encrustation/Scale Light", |
|
"Encrustation/Scale Medium", |
|
"Encrustation/Scale Heavy", |
|
"Debris Grease", |
|
"Debris Silt", |
|
"Roots Fine", |
|
"Roots Mass", |
|
"Weeper", |
|
"Roots Tap", |
|
"Infiltration Light Seeping Joint", |
|
"Infiltration Light Dripping Joint", |
|
"Infiltration Moderate Running Joint", |
|
"Infiltration Severe Gusher Joint", |
|
"Connection Intruding", |
|
"Connection Defective", |
|
"Sealing Ring Intruding", |
|
"Sealing Ring Broken", |
|
"Other sealent intruding", |
|
"Camera Underwater" |
|
] |
|
|
|
|
|
|
|
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", |
|
"RL": "Roots Lateral", |
|
"LR": "Line Right", |
|
"LL": "Line Left", |
|
"LD": "Line Down", |
|
"LU": "Line Up", |
|
"DEF": "Deformed Elliptical (Flexible)", |
|
"WL": "Weeper", |
|
"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 simple_flag(image, label): |
|
|
|
pil_image = Image.fromarray(image.astype(np.uint8)) |
|
img_byte_arr = io.BytesIO() |
|
pil_image.save(img_byte_arr, format='PNG') |
|
|
|
|
|
unique_id = str(uuid.uuid4()) |
|
img_filename = f"{unique_id}.png" |
|
|
|
|
|
image_bytes = img_byte_arr.getvalue() |
|
|
|
|
|
label_dir = f"{label}/{img_filename}" |
|
upload_file( |
|
path_or_fileobj=io.BytesIO(image_bytes), |
|
path_in_repo=label_dir, |
|
repo_id=DATASET_NAME, |
|
repo_type="dataset", |
|
token=HF_TOKEN, |
|
commit_message=f"Add image with label {label}" |
|
) |
|
|
|
return "Thank you for your contribution to the open-source world! Your feedback helps us all move towards a clearer future" |
|
|
|
def replace_label_with_full_name(res, defect_dict_key_code): |
|
new_res = {} |
|
for dic in res: |
|
|
|
parts = dic["label"].split('_', 1) |
|
code = parts[0] |
|
suffix = '_' + parts[1] if len(parts) > 1 else '' |
|
|
|
|
|
full_name = defect_dict_key_code.get(code, code) |
|
|
|
|
|
new_label = full_name + suffix |
|
new_res[new_label] = dic["score"] |
|
|
|
return new_res |
|
|
|
def classify_image(image): |
|
|
|
PIL_image = Image.fromarray(image).convert('RGB') |
|
|
|
|
|
res = pipe(PIL_image) |
|
|
|
|
|
return replace_label_with_full_name(res, defect_dict) |
|
|
|
|
|
|
|
with gr.Blocks() as demo: |
|
gr.Markdown("# Defects Classification with AI by Subterra") |
|
gr.Markdown("Upload an image to view a classification demonstration leveraging the dataset/library of images collected by WRc & United Utilities during The Water Services Regulation Authority (OFWAT) Innovation Challenge – Artificial Intelligence and Sewers. Not only can you see the initial classification, but you as the user can also inform us if the classification is correct. Your response will be used to retrain this model. The team at Subterra would like to thank all of those involved in collecting this dataset as we hope that other groups will use it to further advance technology solutions for the water industry.") |
|
|
|
with gr.Row(): |
|
with gr.Column(): |
|
img_input = gr.Image() |
|
submit_button = gr.Button("Classify") |
|
examples = gr.Examples(["examples/CS.jpg", "examples/GI.jpg", "examples/PP.jpg", "examples/RC.jpg"], label = "Explore Examples", inputs=img_input) |
|
with gr.Column(): |
|
output_label = gr.Label() |
|
flagging_options = gr.Radio(defect_full_names_list, label="Does this classification look off to you? Your sharp eyes can help correct it. Flag any inaccuracies and suggest the right label!") |
|
flag_button = gr.Button("Flag") |
|
flag_status = gr.Textbox(label = "Every flag you submit polishes our dataset. Thanks for being an active participant in our open-source journey.",visible=True) |
|
|
|
submit_button.click(classify_image, inputs=img_input, outputs=output_label) |
|
flag_button.click(simple_flag, inputs=[img_input, flagging_options], outputs=flag_status) |
|
|
|
demo.launch() |
|
|