Upload 2 files

app.py

@@ -0,0 +1,164 @@
+import cv2
+import numpy as np
+import gradio as gr
+
+# ================= FIXED CONFIG =================
+RUST_SCORE_THRESH = 0.45
+
+# GrabCut rectangle (tower prior)
+GC_LEFT   = 0.20
+GC_RIGHT  = 0.80
+GC_TOP    = 0.02
+GC_BOTTOM = 0.98
+
+# Sky HSV bounds
+SKY_LOWER = (80, 20, 40)
+SKY_UPPER = (140, 255, 255)
+
+GROUND_FRAC = 0.15
+GRABCUT_ITERS = 10
+TOWER_MORPH_K = 5
+RUST_MORPH_K = 7
+
+# Rust scoring weights
+RUST_HUE_CENTER = 15
+HUE_WIDTH = 0.20
+W_A, W_H, W_S, W_T, W_L = 0.38, 0.22, 0.15, 0.20, 0.05
+# ================================================
+
+
+def remove_sky_mask(img):
+    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
+    sky = cv2.inRange(hsv, np.array(SKY_LOWER), np.array(SKY_UPPER))
+    return (sky == 0).astype(np.uint8)
+
+
+def remove_ground_mask(img):
+    h, w = img.shape[:2]
+    m = np.ones((h, w), np.uint8)
+    cut = int(h * GROUND_FRAC)
+    m[h - cut:h, :] = 0
+    return m
+
+
+def grabcut_tower_mask(img, non_sky, non_ground):
+    h, w = img.shape[:2]
+
+    mask = np.full((h, w), cv2.GC_PR_BGD, np.uint8)
+    mask[non_sky == 0] = cv2.GC_BGD
+    mask[non_ground == 0] = cv2.GC_BGD
+
+    x1, x2 = int(w * GC_LEFT), int(w * GC_RIGHT)
+    y1, y2 = int(h * GC_TOP), int(h * GC_BOTTOM)
+
+    area = mask[y1:y2, x1:x2]
+    area[non_ground[y1:y2, x1:x2] == 1] = cv2.GC_PR_FGD
+    mask[y1:y2, x1:x2] = area
+
+    bgdModel = np.zeros((1, 65), np.float64)
+    fgdModel = np.zeros((1, 65), np.float64)
+
+    cv2.grabCut(img, mask, None, bgdModel, fgdModel, GRABCUT_ITERS, cv2.GC_INIT_WITH_MASK)
+
+    fg = ((mask == cv2.GC_FGD) | (mask == cv2.GC_PR_FGD)).astype(np.uint8)
+
+    k = np.ones((TOWER_MORPH_K, TOWER_MORPH_K), np.uint8)
+    fg = cv2.morphologyEx(fg, cv2.MORPH_CLOSE, k, iterations=2)
+    fg = cv2.morphologyEx(fg, cv2.MORPH_OPEN, k, iterations=1)
+
+    return fg
+
+
+def rust_score_map(img):
+    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB).astype(np.float32)
+    L, A = lab[:, :, 0], lab[:, :, 1]
+
+    A_n = (A - A.min()) / (A.max() - A.min() + 1e-6)
+    L_n = (L - L.min()) / (L.max() - L.min() + 1e-6)
+
+    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype(np.float32)
+    H, S = hsv[:, :, 0] / 179.0, hsv[:, :, 1] / 255.0
+
+    rust_center = RUST_HUE_CENTER / 179.0
+    hue_score = 1.0 - np.clip(np.abs(H - rust_center) / HUE_WIDTH, 0, 1)
+
+    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    tex = np.abs(cv2.Laplacian(gray, cv2.CV_32F))
+    tex_n = (tex - tex.min()) / (tex.max() - tex.min() + 1e-6)
+
+    score = (
+        W_A * A_n +
+        W_H * hue_score +
+        W_S * S +
+        W_T * tex_n +
+        W_L * (1.0 - np.abs(L_n - 0.55))
+    )
+
+    return np.clip(score, 0, 1)
+
+
+def postprocess(m):
+    k = np.ones((RUST_MORPH_K, RUST_MORPH_K), np.uint8)
+    m = cv2.morphologyEx(m, cv2.MORPH_CLOSE, k, iterations=2)
+    m = cv2.morphologyEx(m, cv2.MORPH_OPEN, k, iterations=1)
+    return m
+
+
+def run(img_rgb):
+    if img_rgb is None:
+        return None, None, None, None, "Upload an image."
+
+    img = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
+
+    non_sky = remove_sky_mask(img)
+    non_ground = remove_ground_mask(img)
+    tower = grabcut_tower_mask(img, non_sky, non_ground)
+
+    roi = tower * non_sky * non_ground
+    score = rust_score_map(img)
+
+    rust = ((score >= RUST_SCORE_THRESH).astype(np.uint8)) * roi
+    rust = postprocess(rust)
+
+    overlay = img.copy()
+    overlay[rust == 1] = [0, 0, 255]
+
+    rust_px = int(rust.sum())
+    roi_px = int(roi.sum())
+    pct = 100 * rust_px / max(1, roi_px)
+
+    stats = (
+        f"Rust threshold: {RUST_SCORE_THRESH}\n"
+        f"Rust pixels: {rust_px}\n"
+        f"ROI pixels: {roi_px}\n"
+        f"Rust on tower ROI: {pct:.2f}%"
+    )
+
+    return (
+        cv2.cvtColor(tower * 255, cv2.COLOR_GRAY2RGB),
+        cv2.cvtColor(rust * 255, cv2.COLOR_GRAY2RGB),
+        cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB),
+        cv2.applyColorMap((score * 255).astype(np.uint8), cv2.COLORMAP_TURBO),
+        stats
+    )
+
+
+with gr.Blocks(title="Tower Rust Detection") as demo:
+    gr.Markdown("## Tower Rust Detection (Fully Automatic)\nUpload an image and click **Run**.")
+
+    inp = gr.Image(type="numpy", label="Input Image")
+    run_btn = gr.Button("Run", variant="primary")
+
+    with gr.Row():
+        out_tower = gr.Image(label="Tower Mask")
+        out_rust = gr.Image(label="Rust Mask")
+    with gr.Row():
+        out_overlay = gr.Image(label="Overlay")
+        out_heat = gr.Image(label="Rust Score Heatmap")
+
+    stats = gr.Textbox(label="Stats", lines=5)
+
+    run_btn.click(run, inputs=inp, outputs=[out_tower, out_rust, out_overlay, out_heat, stats])
+
+if __name__ == "__main__":
+    demo.launch()

best_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8198e6f67a81c4d1f651e58a65c3cba804c8eb1cca6276dbaf65603cb3f6c0c8
+size 97919828