Update app.py

app.py

@@ -1,164 +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()
+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 \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()