Spaces:
Runtime error
Runtime error
| 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() |