import cv2
import numpy as np
from PIL import Image, ImageDraw
import gradio as gr

def classify_pipe_material(image_np):
    """
    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)
    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):
    """
    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)
    
    # 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:
        return "Corrosion"
    else:
        return "Other Defect"

def detect_pipe_issues(image: Image.Image):
    try:
        # Convert PIL image to a NumPy array (RGB)
        image_np = np.array(image)
        annotated = image.copy()  # Copy for annotation
        draw = ImageDraw.Draw(annotated)
        
        # Classify overall pipe material
        pipe_material = classify_pipe_material(image_np)
        
        # 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)
        
        # Apply Gaussian blur to reduce noise
        blurred = cv2.GaussianBlur(enhanced, (5, 5), 0)
        
        # Adaptive thresholding to highlight potential defect areas
        thresh = cv2.adaptiveThreshold(
            blurred, 255,
            cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
            cv2.THRESH_BINARY_INV,
            11, 2
        )
        
        # Morphological closing to connect fragmented regions
        kernel = np.ones((3, 3), np.uint8)
        morph = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel, iterations=2)
        
        # Edge detection
        edges = cv2.Canny(morph, 50, 150)
        
        # 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)
            # 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 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)
        
        # Create a summary including pipe material and detected defects
        if detections:
            summary = f"Pipe Material: {pipe_material}\nDetected Issues:\n" + "\n".join(detections)
        else:
            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}"

iface = gr.Interface(
    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="Pipe 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."
    )
)

if __name__ == "__main__":
    iface.launch()