app.py

import os
import cv2
import torch
import numpy as np
import gradio as gr
import matplotlib.pyplot as plt
import pandas as pd
from glob import glob
from PIL import Image
from skimage.measure import regionprops, label
from scipy.spatial.distance import cdist
from scipy.spatial import Delaunay
from io import BytesIO
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
import segmentation_models_pytorch as smp

# Configuration
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DIAMETER_MM = 152.4
MIN_SIZE = 256

class PetModel(torch.nn.Module):
    def __init__(self, arch, encoder_name, in_channels, out_classes, **kwargs):
        super().__init__()
        self.model = smp.create_model(
            arch, encoder_name, in_channels=in_channels, classes=out_classes, **kwargs
        )
        params = smp.encoders.get_preprocessing_params(encoder_name)
        self.register_buffer("std", torch.tensor(params["std"]).view(1, 3, 1, 1))
        self.register_buffer("mean", torch.tensor(params["mean"]).view(1, 3, 1, 1))

    def forward(self, image):
        image = (image - self.mean) / self.std
        return self.model(image)

def preprocess_image(image, min_size=MIN_SIZE):
    image = np.array(image)
    if len(image.shape) == 2:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    elif image.shape[2] == 4:
        image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
    elif image.shape[2] == 1:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)

    original_size = image.shape[:2]
    h, w = image.shape[:2]
    if h < min_size or w < min_size:
        new_size = (max(w, min_size), max(h, min_size))
        image = cv2.resize(image, new_size, interpolation=cv2.INTER_LINEAR)

    image = image.astype(np.float32) / 255.0
    image = torch.tensor(image).permute(2, 0, 1).unsqueeze(0)
    return image, original_size

def postprocess_output(output, original_size):
    prob_mask = output.sigmoid()
    pred_mask = (prob_mask > 0.5).float()
    pred_mask = pred_mask.squeeze().cpu().numpy()
    if pred_mask.shape != original_size:
        pred_mask = cv2.resize(pred_mask, (original_size[1], original_size[0]), interpolation=cv2.INTER_NEAREST)
    return pred_mask

def load_model(model_path):
    model = PetModel("unet", "efficientnet-b5", in_channels=3, out_classes=1)
    model.load_state_dict(torch.load(model_path, map_location=DEVICE))
    model = model.to(DEVICE)
    model.eval()
    return model

model = load_model("segmentation_model_final.pth")

csv_output_path = "measurement_summary.csv"

def fig_to_image(fig):
    buf = BytesIO()
    canvas = FigureCanvas(fig)
    canvas.print_png(buf)
    buf.seek(0)
    return Image.open(buf)

def analyze(image):
    input_tensor, original_size = preprocess_image(image)
    input_tensor = input_tensor.to(DEVICE)

    with torch.no_grad():
        output = model(input_tensor)

    prediction_mask = postprocess_output(output, original_size)
    image_np = np.array(image)
    gray_img = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
    label_img = (prediction_mask * 255).astype(np.uint8)

    _, bw = cv2.threshold(gray_img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    contours, _ = cv2.findContours(bw, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
    contours = sorted(contours, key=cv2.contourArea, reverse=True)
    if not contours:
        return None, None, None, "No contour found."

    boundary = contours[0].squeeze()
    dist_matrix = cdist(boundary, boundary)
    i, j = np.unravel_index(np.argmax(dist_matrix), dist_matrix.shape)
    line_pts = np.array([boundary[i], boundary[j]])
    pixel_diameter = np.linalg.norm(boundary[i] - boundary[j])
    pixels_per_mm = pixel_diameter / DIAMETER_MM
    pixel_length_mm = 1 / pixels_per_mm
    line_length_mm = pixel_diameter * pixel_length_mm

    fig1 = plt.figure(figsize=(6, 6))
    plt.imshow(image_np)
    plt.plot(boundary[:, 0], boundary[:, 1], 'g', linewidth=2)
    plt.plot(line_pts[:, 0], line_pts[:, 1], 'r', linewidth=2)
    plt.title(f"Calibration Line: {line_length_mm:.2f} mm")
    plt.axis("off")
    img1 = fig_to_image(fig1)

    binary_mask = (label_img > 127).astype(np.uint8)
    color_mask = cv2.cvtColor(label_img, cv2.COLOR_GRAY2BGR)

    feret_lengths, feret_widths, rectangles = [], [], []
    contours_mask, _ = cv2.findContours(binary_mask * 255, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    for cnt in contours_mask:
        if len(cnt) >= 5:
            rect = cv2.minAreaRect(cnt)
            box = cv2.boxPoints(rect).astype(np.intp)
            width, height = rect[1]
            feret_length = max(width, height)
            feret_lengths.append(feret_length)
            feret_widths.append(min(width, height))
            rectangles.append((box, feret_length))

    thresholds = np.percentile(feret_lengths, [20, 40, 60, 80]) if feret_lengths else [0]*4
    colors = [(0,0,255),(0,128,255),(0,255,255),(0,255,0),(255,0,0)]
    for box, length in rectangles:
        if length <= thresholds[0]: color = colors[0]
        elif length <= thresholds[1]: color = colors[1]
        elif length <= thresholds[2]: color = colors[2]
        elif length <= thresholds[3]: color = colors[3]
        else: color = colors[4]
        cv2.drawContours(color_mask, [box], 0, color, 8)

    fig2 = plt.figure(figsize=(6, 6))
    plt.imshow(cv2.cvtColor(color_mask, cv2.COLOR_BGR2RGB))
    plt.title("Feret Rectangles (Colored by Size)")
    plt.axis("off")
    img2 = fig_to_image(fig2)

    labeled_img = label(binary_mask)
    props = regionprops(labeled_img)
    centroids = np.array([p.centroid for p in props])

    edge_lengths = []
    fig3 = plt.figure(figsize=(6, 6))
    plt.imshow(label_img, cmap="gray")
    if len(centroids) >= 3:
        tri = Delaunay(centroids)
        plt.triplot(centroids[:, 1], centroids[:, 0], tri.simplices.copy(), color="red", linewidth=1)
        for simplex in tri.simplices:
            for i in range(3):
                pt1 = centroids[simplex[i]]
                pt2 = centroids[simplex[(i + 1) % 3]]
                dist_px = np.linalg.norm(pt1 - pt2)
                dist_mm = dist_px * pixel_length_mm
                edge_lengths.append(dist_mm)
        plt.title("Delaunay Triangulation")
    else:
        plt.title("Not Enough Aggregates for Triangulation")
    plt.axis("off")
    img3 = fig_to_image(fig3)

    num_white_pixels = np.sum(binary_mask == 1)
    num_nonblack_pixels = np.count_nonzero(gray_img)
    aggregate_area_mm2 = num_white_pixels * (pixel_length_mm ** 2)
    total_area_mm2 = num_nonblack_pixels * (pixel_length_mm ** 2)
    aggregate_ratio = aggregate_area_mm2 / total_area_mm2 if total_area_mm2 > 0 else 0

    if feret_lengths:
        avg_feret_length_mm = np.mean(feret_lengths) * pixel_length_mm
        avg_feret_width_mm = np.mean(feret_widths) * pixel_length_mm
        max_feret_length_mm = np.max(feret_lengths) * pixel_length_mm
        roundness_aggregate = avg_feret_length_mm / avg_feret_width_mm
    else:
        avg_feret_length_mm = avg_feret_width_mm = max_feret_length_mm = roundness_aggregate = 0

    # Save to CSV
    data = {
        "Pixel_Size_mm_per_pixel": [pixel_length_mm],
        "Aggregate_Area_mm2": [aggregate_area_mm2],
        "Aggregate_Ratio": [aggregate_ratio],
        "Avg_Length_mm": [avg_feret_length_mm],
        "Avg_Width_mm": [avg_feret_width_mm],
        "Max_Length_mm": [max_feret_length_mm],
        "Roundness": [roundness_aggregate],
        "Avg_Dist_mm": [np.mean(edge_lengths) if edge_lengths else 0],
        "Max_Dist_mm": [np.max(edge_lengths) if edge_lengths else 0]
    }
    df = pd.DataFrame(data)
    df.to_csv(csv_output_path, index=False)

    summary = f"""📏 **Measurements Summary**:
- Pixel Size: `{pixel_length_mm:.4f}` mm/pixel
- Aggregate Area: `{aggregate_area_mm2:.2f}` mm²
- Aggregate Ratio: `{aggregate_ratio:.4f}`
- Avg Aggregate Length: `{avg_feret_length_mm:.2f}` mm
- Avg Aggregate Width: `{avg_feret_width_mm:.2f}` mm
- Max Aggregate Length: `{max_feret_length_mm:.2f}` mm
- Aggregate Roundness: `{roundness_aggregate:.2f}`
"""
    if edge_lengths:
        summary += f"- Avg Inter-Aggregate Distance: `{np.mean(edge_lengths):.2f}` mm\n"
        summary += f"- Max Inter-Aggregate Distance: `{np.max(edge_lengths):.2f}` mm\n"

    return img1, img2, img3, summary

demo = gr.Interface(
    fn=analyze,
    inputs=gr.Image(type="pil", label="Upload Concrete Image"),
    outputs=[
        gr.Image(label="Boundary & Calibration Line"),
        gr.Image(label="Feret Rectangles"),
        gr.Image(label="Delaunay Triangulation"),
        gr.Textbox(label="Measurements Summary")
    ],
    title="Concrete Aggregate Analysis App",
    description="Upload a concrete image. The model will segment aggregates and analyze their distribution and shape."
)

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