Update app.py

app.py

@@ -5,7 +5,7 @@ import gradio as gr
 
 def classify_pipe_material(image_np):
     """
-    Heuristic to classify the overall pipe material based on brightness.
+    Classify overall material based on image brightness.
     Brighter images (mean intensity > 130) are assumed to be Plastic; otherwise, Metal.
     """
     gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
@@ -14,36 +14,62 @@ def classify_pipe_material(image_np):
 
 def detect_rust(roi):
     """
-    Detect rust in the region of interest (ROI) by analyzing the HSV color space.
-    Rust typically has reddish-brown hues.
+    Detect rust by checking for reddish-brown hues in the ROI.
     """
-    # Convert ROI to HSV color space
     hsv_roi = cv2.cvtColor(roi, cv2.COLOR_RGB2HSV)
-    # Define rust color range in HSV (tweak these values as needed)
     lower_rust = np.array([5, 50, 50])
     upper_rust = np.array([25, 255, 255])
     mask = cv2.inRange(hsv_roi, lower_rust, upper_rust)
     rust_ratio = np.count_nonzero(mask) / float(roi.shape[0] * roi.shape[1])
     return rust_ratio
 
+def detect_mold(roi):
+    """
+    Detect mold by looking for greenish hues, which may indicate fungal growth.
+    """
+    hsv_roi = cv2.cvtColor(roi, cv2.COLOR_RGB2HSV)
+    lower_mold = np.array([35, 50, 20])
+    upper_mold = np.array([85, 255, 120])
+    mask = cv2.inRange(hsv_roi, lower_mold, upper_mold)
+    mold_ratio = np.count_nonzero(mask) / float(roi.shape[0] * roi.shape[1])
+    return mold_ratio
+
+def detect_water_damage(roi):
+    """
+    Detect water damage by checking for discoloration typical of stains (dark brownish-yellow).
+    """
+    hsv_roi = cv2.cvtColor(roi, cv2.COLOR_RGB2HSV)
+    lower_water = np.array([5, 50, 50])
+    upper_water = np.array([20, 200, 150])
+    mask = cv2.inRange(hsv_roi, lower_water, upper_water)
+    water_ratio = np.count_nonzero(mask) / float(roi.shape[0] * roi.shape[1])
+    return water_ratio
+
 def classify_defect(roi):
     """
-    Classify the defect type using both geometric/texture heuristics and color analysis.
-    The function returns one of:
-      - "Rust" (if a significant fraction of the region has rust-like colors)
-      - "Crack" (if the ROI is small, long, and has high intensity variation)
-      - "Corrosion" (if the ROI is larger with moderate texture variation)
-      - "Other Defect" (fallback category)
+    Classify the defect using a combination of color and texture heuristics.
+    Priority is given to color cues:
+      - "Rust": reddish-brown
+      - "Mold": greenish
+      - "Water Damage": discoloration from water stains
+    Then geometric/texture analysis is used to differentiate "Crack" and "Corrosion."
     """
     area = roi.shape[0] * roi.shape[1]
     std_intensity = np.std(roi)
     
-    # Check for rust first
     rust_ratio = detect_rust(roi)
-    if rust_ratio > 0.3:
+    mold_ratio = detect_mold(roi)
+    water_ratio = detect_water_damage(roi)
+    
+    # Check for color-based defects first.
+    if rust_ratio > 0.25:
         return "Rust"
+    elif mold_ratio > 0.2:
+        return "Mold"
+    elif water_ratio > 0.2:
+        return "Water Damage"
     
-    # Use area and intensity variation to distinguish other defects.
+    # Then use size and intensity variation for structural defects.
     if area < 5000 and std_intensity > 50:
         return "Crack"
     elif area >= 5000 and std_intensity > 40:
@@ -51,25 +77,23 @@ def classify_defect(roi):
     else:
         return "Other Defect"
 
-def detect_pipe_issues(image: Image.Image):
+def detect_infrastructure_issues(image: Image.Image):
     try:
-        # Convert PIL image to a NumPy array (RGB)
+        # Convert the PIL image to a NumPy array (RGB)
         image_np = np.array(image)
-        annotated = image.copy()  # Copy for annotation
+        annotated = image.copy()  # For drawing annotations
         draw = ImageDraw.Draw(annotated)
         
-        # Classify overall pipe material
-        pipe_material = classify_pipe_material(image_np)
+        # Classify the overall pipe (or structure) material
+        overall_material = classify_pipe_material(image_np)
         
-        # Preprocessing: convert to grayscale and enhance contrast with CLAHE
+        # Enhance image for defect detection:
         gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
         clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
         enhanced = clahe.apply(gray)
-        
-        # Apply Gaussian blur to reduce noise
         blurred = cv2.GaussianBlur(enhanced, (5, 5), 0)
         
-        # Adaptive thresholding to highlight potential defect areas
+        # Adaptive thresholding highlights potential defect areas
         thresh = cv2.adaptiveThreshold(
             blurred, 255,
             cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
@@ -77,31 +101,31 @@ def detect_pipe_issues(image: Image.Image):
             11, 2
         )
         
-        # Morphological closing to connect fragmented regions
+        # Morphological closing to consolidate defect regions
         kernel = np.ones((3, 3), np.uint8)
         morph = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=2)
         
-        # Edge detection
+        # Edge detection to identify defect boundaries
         edges = cv2.Canny(morph, 50, 150)
         
-        # Find contours corresponding to potential defects
+        # Extract contours as candidate defect regions
         contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
         detections = []
         
-        # Define colors for different defect types
+        # Define distinct colors for each defect type
         colors = {
             "Rust": "orange",
+            "Mold": "purple",
+            "Water Damage": "blue",
             "Crack": "red",
-            "Corrosion": "blue",
-            "Other Defect": "green"
+            "Corrosion": "cyan",
+            "Other Defect": "gray"
         }
         
         for cnt in contours:
-            # Filter out small contours to ignore noise
-            if cv2.contourArea(cnt) < 100:
+            if cv2.contourArea(cnt) < 100:  # Filter out noise
                 continue
             x, y, w, h = cv2.boundingRect(cnt)
-            # Extract ROI from the original image
             roi = image_np[y:y+h, x:x+w]
             if roi.size == 0:
                 continue
@@ -109,16 +133,16 @@ def detect_pipe_issues(image: Image.Image):
             detection_info = f"{defect_type} at ({x}, {y}, {w}, {h})"
             detections.append(detection_info)
             
-            # Draw bounding box with corresponding color and label
-            box_color = colors.get(defect_type, "green")
+            # Draw bounding box with defect-specific color
+            box_color = colors.get(defect_type, "gray")
             draw.rectangle([x, y, x+w, y+h], outline=box_color, width=2)
             draw.text((x, y-10), defect_type, fill=box_color)
         
-        # Create a summary including pipe material and detected defects
+        # Create a textual summary
         if detections:
-            summary = f"Pipe Material: {pipe_material}\nDetected Issues:\n" + "\n".join(detections)
+            summary = f"Overall Material: {overall_material}\nDetected Issues:\n" + "\n".join(detections)
         else:
-            summary = f"Pipe Material: {pipe_material}\nNo significant defects detected."
+            summary = f"Overall Material: {overall_material}\nNo significant defects detected."
         
         return annotated, summary
     except Exception as e:
@@ -126,14 +150,14 @@ def detect_pipe_issues(image: Image.Image):
         return image, f"Error: {e}"
 
 iface = gr.Interface(
-    fn=detect_pipe_issues,
-    inputs=gr.Image(type="pil", label="Upload a Pipe Image"),
+    fn=detect_infrastructure_issues,
+    inputs=gr.Image(type="pil", label="Upload an Infrastructure Image"),
     outputs=[gr.Image(label="Annotated Image"), gr.Textbox(label="Detection Summary")],
-    title="Pipe Defect Detector",
+    title="Comprehensive Home Infrastructure Defect Detector",
     description=(
-        "Upload an image of a pipe to detect granular issues such as cracks, corrosion, rust, "
-        "and other defects. The app classifies the defect type and displays a colored bounding box for each class. "
-        "Pipe material (Plastic or Metal) is also identified."
+        "Upload an image of a pipe or any home infrastructure (walls, floors, etc.) to detect defects. "
+        "This tool identifies issues such as Rust (orange), Mold (purple), Water Damage (blue), Cracks (red), "
+        "and Corrosion (cyan), and returns both an annotated image and a detailed summary."
     )
 )