Update app.py

app.py

@@ -1,23 +1,54 @@
 import cv2
 import numpy as np
-from PIL import Image, ImageDraw, ImageFont
+from PIL import Image, ImageDraw
 import gradio as gr
 
-def detect_cracks(image: Image.Image):
+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.
+    """
+    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 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).
+    """
+    area = roi.shape[0] * roi.shape[1]
+    std_intensity = np.std(roi)
+    # Adjust thresholds based on real-world testing.
+    if area < 5000 and std_intensity > 50:
+        return "Crack"
+    elif area >= 5000 and std_intensity > 40:
+        return "Corrosion"
+    else:
+        return "Other Defect"
+
+def detect_pipe_issues(image: Image.Image):
     try:
-        # Convert the PIL image to an OpenCV (RGB) image
-        rgb_image = np.array(image)
-        # Also create a copy for annotation
-        annotated = image.copy()
+        # Convert the PIL image to a NumPy array (RGB)
+        image_np = np.array(image)
+        annotated = image.copy()  # Copy for annotation
         draw = ImageDraw.Draw(annotated)
         
-        # Convert to grayscale
-        gray = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2GRAY)
+        # Classify the overall pipe material
+        pipe_material = classify_pipe_material(image_np)
+        
+        # Convert to grayscale and enhance contrast using CLAHE
+        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(gray, (5, 5), 0)
+        blurred = cv2.GaussianBlur(enhanced, (5, 5), 0)
         
-        # Adaptive thresholding to highlight crack-like features
+        # Adaptive thresholding to highlight potential defect areas
         thresh = cv2.adaptiveThreshold(
             blurred, 255,
             cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
@@ -25,70 +56,50 @@ def detect_cracks(image: Image.Image):
             11, 2
         )
         
-        # Use morphological closing to fill gaps in potential cracks
+        # Use morphological closing to connect fragmented areas
         kernel = np.ones((3, 3), np.uint8)
         morph = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=2)
         
-        # Edge detection
+        # Detect edges with Canny
         edges = cv2.Canny(morph, 50, 150)
         
-        # Find contours from edges
+        # Find contours that might correspond to defects
         contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        
-        detections = []  # to hold detection details
+        detections = []
         
         for cnt in contours:
-            # Filter out small contours (noise)
-            if cv2.arcLength(cnt, True) > 100:
-                x, y, w, h = cv2.boundingRect(cnt)
-                
-                # Extract ROI from the original image (for material classification)
-                roi = rgb_image[y:y+h, x:x+w]
-                if roi.size == 0:
-                    continue
-                # Convert ROI to grayscale and compute mean intensity
-                roi_gray = cv2.cvtColor(roi, cv2.COLOR_RGB2GRAY)
-                mean_intensity = np.mean(roi_gray)
-                
-                # Simple heuristic: classify material based on brightness
-                # (These thresholds are arbitrary and should be tuned based on real data.)
-                if mean_intensity < 80:
-                    material = "Concrete"
-                elif mean_intensity < 150:
-                    material = "Tile"
-                else:
-                    material = "Wood"
-                
-                label = f"Crack ({material})"
-                detections.append(f"Detected crack at ({x}, {y}, {w}, {h}) on {material} (mean intensity: {mean_intensity:.1f})")
-                
-                # Draw rectangle and label on the annotated image
-                draw.rectangle([x, y, x+w, y+h], outline="red", width=2)
-                # Draw the label above the rectangle
-                draw.text((x, y-10), label, fill="red")
+            # 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]
+            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")
         
-        # Create a text summary of detections
+        # Prepare a textual summary including pipe material and all defect detections
         if detections:
-            summary = "\n".join(detections)
+            summary = f"Pipe Material: {pipe_material}\nDetected Issues:\n" + "\n".join(detections)
         else:
-            summary = "No significant cracks detected."
+            summary = f"Pipe Material: {pipe_material}\nNo significant defects detected."
         
         return annotated, summary
-        
     except Exception as e:
         print("Error during detection:", e)
         return image, f"Error: {e}"
 
-# Create a Gradio interface with two outputs: image and text
 iface = gr.Interface(
-    fn=detect_cracks,
-    inputs=gr.Image(type="pil", label="Upload an Image (Floor/Wall)"),
+    fn=detect_pipe_issues,
+    inputs=gr.Image(type="pil", label="Upload a Pipe Image"),
     outputs=[gr.Image(label="Annotated Image"), gr.Textbox(label="Detection Summary")],
-    title="Home Inspection: Granular Crack & Material Detector",
+    title="Pipe Defect Detector",
     description=(
-        "Upload an image of a floor or wall to detect cracks and infer the underlying material "
-        "(Concrete, Tile, or Wood) using classical computer vision techniques. "
-        "This demo returns both an annotated image and a textual summary."
+        "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."
     )
 )