Update app.py

app.py

@@ -5,24 +5,45 @@ import gradio as gr
 
 def classify_pipe_material(image_np):
     """
-    Heuristic to classify the overall pipe material based on the image's brightness.
-    Assumes that plastic pipes tend to be brighter and metal pipes darker.
+    Heuristic to classify the overall pipe material based on brightness.
+    Brighter images (mean intensity > 130) are assumed to be Plastic; otherwise, Metal.
     """
     gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
     mean_intensity = np.mean(gray)
-    # If the mean intensity is high, assume plastic; otherwise, metal.
     return "Plastic" if mean_intensity > 130 else "Metal"
 
+def detect_rust(roi):
+    """
+    Detect rust in the region of interest (ROI) by analyzing the HSV color space.
+    Rust typically has reddish-brown hues.
+    """
+    # 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 classify_defect(roi):
     """
-    Heuristic to classify the defect type within a region of interest (ROI).
-    This simple method uses the area and intensity variation (std deviation)
-    to differentiate between a crack (typically long and thin with high contrast)
-    and corrosion (possibly larger, more blob-like with different texture).
+    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)
     """
     area = roi.shape[0] * roi.shape[1]
     std_intensity = np.std(roi)
-    # Adjust thresholds based on real-world testing.
+    
+    # Check for rust first
+    rust_ratio = detect_rust(roi)
+    if rust_ratio > 0.3:
+        return "Rust"
+    
+    # Use area and intensity variation to distinguish other defects.
     if area < 5000 and std_intensity > 50:
         return "Crack"
     elif area >= 5000 and std_intensity > 40:
@@ -32,15 +53,15 @@ def classify_defect(roi):
 
 def detect_pipe_issues(image: Image.Image):
     try:
-        # Convert the PIL image to a NumPy array (RGB)
+        # Convert PIL image to a NumPy array (RGB)
         image_np = np.array(image)
         annotated = image.copy()  # Copy for annotation
         draw = ImageDraw.Draw(annotated)
         
-        # Classify the overall pipe material
+        # Classify overall pipe material
         pipe_material = classify_pipe_material(image_np)
         
-        # Convert to grayscale and enhance contrast using CLAHE
+        # Preprocessing: convert to grayscale and enhance contrast with CLAHE
         gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
         clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
         enhanced = clahe.apply(gray)
@@ -56,32 +77,44 @@ def detect_pipe_issues(image: Image.Image):
             11, 2
         )
         
-        # Use morphological closing to connect fragmented areas
+        # Morphological closing to connect fragmented regions
         kernel = np.ones((3, 3), np.uint8)
         morph = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=2)
         
-        # Detect edges with Canny
+        # Edge detection
         edges = cv2.Canny(morph, 50, 150)
         
-        # Find contours that might correspond to defects
+        # Find contours corresponding to potential defects
         contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
         detections = []
         
+        # Define colors for different defect types
+        colors = {
+            "Rust": "orange",
+            "Crack": "red",
+            "Corrosion": "blue",
+            "Other Defect": "green"
+        }
+        
         for cnt in contours:
             # Filter out small contours to ignore noise
             if cv2.contourArea(cnt) < 100:
                 continue
             x, y, w, h = cv2.boundingRect(cnt)
-            roi = enhanced[y:y+h, x:x+w]
+            # Extract ROI from the original image
+            roi = image_np[y:y+h, x:x+w]
+            if roi.size == 0:
+                continue
             defect_type = classify_defect(roi)
             detection_info = f"{defect_type} at ({x}, {y}, {w}, {h})"
             detections.append(detection_info)
             
-            # Draw bounding box and defect label
-            draw.rectangle([x, y, x+w, y+h], outline="red", width=2)
-            draw.text((x, y-10), defect_type, fill="red")
+            # Draw bounding box with corresponding color and label
+            box_color = colors.get(defect_type, "green")
+            draw.rectangle([x, y, x+w, y+h], outline=box_color, width=2)
+            draw.text((x, y-10), defect_type, fill=box_color)
         
-        # Prepare a textual summary including pipe material and all defect detections
+        # Create a summary including pipe material and detected defects
         if detections:
             summary = f"Pipe Material: {pipe_material}\nDetected Issues:\n" + "\n".join(detections)
         else:
@@ -98,8 +131,9 @@ iface = gr.Interface(
     outputs=[gr.Image(label="Annotated Image"), gr.Textbox(label="Detection Summary")],
     title="Pipe Defect Detector",
     description=(
-        "Upload an image of a pipe to detect granular issues such as cracks or corrosion. "
-        "The app classifies the defect type and determines if the pipe material is plastic or metal."
+        "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."
     )
 )