dmncjsph/roof-rust-system

1from pathlib import Path2import io3import os4 5from flask import Flask, request, jsonify, render_template6import torch7import torch.nn as nn8from PIL import Image, UnidentifiedImageError9from torchvision.models import efficientnet_b2, EfficientNet_B2_Weights10 11BASE_DIR = Path(__file__).resolve().parent12 13ROOF_FILTER_MODEL_PATH = BASE_DIR / "efficientnet_b2_roof_filter_best.pth"14RUST_MODEL_PATH = BASE_DIR / "efficientnet_b2_roofrust_best_b2_mixup.pth"15 16app = Flask(17    __name__,18    template_folder=str(BASE_DIR),19    static_folder=str(BASE_DIR),20    static_url_path=""21)22 23device = torch.device("cuda" if torch.cuda.is_available() else "cpu")24print("Using device:", device)25 26ALLOWED_EXTENSIONS = {"jpg", "jpeg", "png", "bmp", "webp"}27 28ROOF_THRESHOLD = 70.029LOW_CONFIDENCE_THRESHOLD = 70.030REJECT_THRESHOLD = 40.031 32DISPLAY_ORDER = [33    "No Rust",34    "Slightly Visible Rust",35    "Visible Rust",36    "Heavy Visible Rust"37]38 39CLASS_DETAILS = {40    "No Rust": {41        "iso_basis": "Closest to Grade A, where the steel surface is largely covered with adhering mill scale and has little or no rust.",42        "recommendation": [43            "Your roof looks in good condition.",44            "No immediate repair is needed.",45            "Keep the roof clean and dry.",46            "Check it regularly, especially after strong rain or storms.",47            "Apply protective coating during scheduled maintenance to help prevent future rust."48        ]49    },50    "Slightly Visible Rust": {51        "iso_basis": "Closest to Grade B, where rusting has begun and mill scale has started to flake.",52        "recommendation": [53            "Early signs of rust are starting to appear.",54            "Clean the affected area as soon as possible.",55            "Remove dirt, loose rust, and debris.",56            "Apply primer or protective paint to stop the rust from spreading.",57            "Monitor the area regularly to make sure the damage does not get worse."58        ]59    },60    "Visible Rust": {61        "iso_basis": "Closest to Grade C, where mill scale has rusted away or can be scraped off and slight pitting is visible.",62        "recommendation": [63            "Rust is already noticeable and should not be ignored.",64            "Schedule maintenance soon to avoid bigger damage.",65            "Clean and prepare the surface properly before repainting or recoating.",66            "Check nearby screws, joints, and fasteners because rust often spreads around these areas.",67            "If the rust covers a larger area, ask a roofing professional to inspect it."68        ]69    },70    "Heavy Visible Rust": {71        "iso_basis": "Closest to Grade D, where mill scale has rusted away and general pitting is visible.",72        "recommendation": [73            "The roof shows serious rust and may already be weakening.",74            "Urgent repair is recommended.",75            "Have the roof inspected for deep corrosion, holes, or possible structural damage.",76            "Heavy rust should be removed thoroughly before any recoating is done.",77            "If the damage is severe, some roof parts may need professional restoration or replacement."78        ]79    }80}81 82def allowed_file(filename):83    return "." in filename and filename.rsplit(".", 1)[1].lower() in ALLOWED_EXTENSIONS84 85if not ROOF_FILTER_MODEL_PATH.exists():86    raise FileNotFoundError(f"Roof filter checkpoint not found: {ROOF_FILTER_MODEL_PATH}")87 88if not RUST_MODEL_PATH.exists():89    raise FileNotFoundError(f"Rust classifier checkpoint not found: {RUST_MODEL_PATH}")90 91preprocess = EfficientNet_B2_Weights.DEFAULT.transforms()92 93roof_checkpoint = torch.load(ROOF_FILTER_MODEL_PATH, map_location=device, weights_only=False)94roof_class_names = roof_checkpoint["class_names"]95roof_num_classes = len(roof_class_names)96 97roof_model = efficientnet_b2(weights=None)98roof_in_features = roof_model.classifier[1].in_features99roof_model.classifier[1] = nn.Linear(roof_in_features, roof_num_classes)100roof_model.load_state_dict(roof_checkpoint["model_state_dict"])101roof_model = roof_model.to(device)102roof_model.eval()103 104rust_checkpoint = torch.load(RUST_MODEL_PATH, map_location=device, weights_only=False)105rust_class_names = rust_checkpoint["class_names"]106rust_num_classes = len(rust_class_names)107 108rust_model = efficientnet_b2(weights=None)109rust_in_features = rust_model.classifier[1].in_features110rust_model.classifier[1] = nn.Linear(rust_in_features, rust_num_classes)111rust_model.load_state_dict(rust_checkpoint["model_state_dict"])112rust_model = rust_model.to(device)113rust_model.eval()114 115def run_roof_filter(image: Image.Image):116    image = image.convert("RGB")117    input_tensor = preprocess(image).unsqueeze(0).to(device)118 119    with torch.no_grad():120        outputs = roof_model(input_tensor)121        probs = torch.softmax(outputs[0], dim=0)122 123    result = {roof_class_names[i]: float(probs[i]) for i in range(roof_num_classes)}124    predicted_class = max(result, key=result.get)125    confidence = result[predicted_class] * 100126    return predicted_class, confidence, result127 128def run_rust_classifier(image: Image.Image):129    image = image.convert("RGB")130    input_tensor = preprocess(image).unsqueeze(0).to(device)131 132    with torch.no_grad():133        outputs = rust_model(input_tensor)134        probs = torch.softmax(outputs[0], dim=0)135 136    raw_result = {rust_class_names[i]: float(probs[i]) for i in range(rust_num_classes)}137    final_class = max(raw_result, key=raw_result.get)138    confidence = raw_result[final_class] * 100139 140    ordered_result = {cls: raw_result.get(cls, 0.0) for cls in DISPLAY_ORDER}141 142    if confidence < REJECT_THRESHOLD:143        return {144            "accepted": False,145            "low_confidence": False,146            "error": "The image is a roof, but the rust severity result is too uncertain. Please upload a clearer roof image."147        }148 149    if confidence < LOW_CONFIDENCE_THRESHOLD:150        return {151            "accepted": True,152            "low_confidence": True,153            "warning": "The image was detected as a roof, but the rust severity prediction is not highly confident. Please review the result carefully or try another clearer image.",154            "predictions": ordered_result,155            "final_classification": final_class,156            "confidence": round(confidence, 2),157            "iso_basis": CLASS_DETAILS[final_class]["iso_basis"],158            "recommendation": CLASS_DETAILS[final_class]["recommendation"]159        }160 161    return {162        "accepted": True,163        "low_confidence": False,164        "warning": "",165        "predictions": ordered_result,166        "final_classification": final_class,167        "confidence": round(confidence, 2),168        "iso_basis": CLASS_DETAILS[final_class]["iso_basis"],169        "recommendation": CLASS_DETAILS[final_class]["recommendation"]170    }171 172@app.route("/")173def home():174    return render_template("home.html")175 176@app.route("/analyzer")177def analyzer():178    return render_template("analyzer.html")179 180@app.route("/about-team")181def about_team():182    return render_template("about_team.html")183 184@app.route("/predict", methods=["POST"])185def predict():186    if "file" not in request.files:187        return jsonify({"error": "No file uploaded"}), 400188 189    file = request.files["file"]190 191    if file.filename == "":192        return jsonify({"error": "Empty filename"}), 400193 194    if not allowed_file(file.filename):195        return jsonify({196            "error": "Unsupported file type. Please upload JPG, JPEG, PNG, BMP, or WEBP."197        }), 400198 199    try:200        image_bytes = file.read()201 202        if not image_bytes:203            return jsonify({"error": "Uploaded file is empty"}), 400204 205        image_stream = io.BytesIO(image_bytes)206 207        test_img = Image.open(image_stream)208        test_img.verify()209 210        image_stream.seek(0)211        image = Image.open(image_stream).convert("RGB")212 213        roof_class, roof_confidence, _ = run_roof_filter(image)214 215        roof_label_map = {name.lower(): name for name in roof_class_names}216        predicted_roof_label = roof_label_map.get("roof", "roof")217 218        if roof_class != predicted_roof_label or roof_confidence < ROOF_THRESHOLD:219            return jsonify({220                "error": "Please upload a roof image only. The uploaded image was detected as non-roof or too uncertain.",221                "roof_filter_prediction": roof_class,222                "roof_filter_confidence": round(roof_confidence, 2)223            }), 400224 225        rust_result = run_rust_classifier(image)226 227        if not rust_result["accepted"]:228            return jsonify({229                "error": rust_result["error"],230                "roof_filter_prediction": roof_class,231                "roof_filter_confidence": round(roof_confidence, 2)232            }), 400233 234        rust_result["roof_filter_prediction"] = roof_class235        rust_result["roof_filter_confidence"] = round(roof_confidence, 2)236 237        return jsonify(rust_result)238 239    except UnidentifiedImageError:240        return jsonify({241            "error": "Invalid or corrupted image file. Please upload a valid image."242        }), 400243 244    except Exception as e:245        return jsonify({"error": str(e)}), 500246 247if __name__ == "__main__":248    port = int(os.environ.get("PORT", 7860))249    app.run(host="0.0.0.0", port=port, debug=False)