Update app.py

app.py

@@ -1,221 +1,164 @@
-# Hugging Face Gradio App: Aggregate Analysis with Feret and Boundary Extension
-
-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")
-
-def fig_to_image(fig):
-    buf = BytesIO()
-    canvas = FigureCanvas(fig)
-    canvas.print_png(buf)
-    buf.seek(0)
-    return Image.open(buf)
-
-def predict(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)
-
-    # Calibration using outer boundary
-    _, bw = cv2.threshold(gray_img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
-    contours, _ = cv2.findContours(bw, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
-    pixel_length_mm = 1.0
-    dia = DIAMETER_MM
-    summary = {}
-
-    figs = []
-    if contours:
-        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 / dia
-        pixel_length_mm = 1 / pixels_per_mm
-        line_length_mm = pixel_diameter * pixel_length_mm
-
-        # Boundary Plot
-        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'Line Length: {line_length_mm:.2f} mm')
-        plt.axis('off')
-        figs.append(fig1)
-
-    # Feret Analysis
-    label_img = (prediction_mask > 0.5).astype(np.uint8)
-    binary_mask = (label_img * 255).astype(np.uint8)
-    color_mask = cv2.cvtColor(binary_mask, cv2.COLOR_GRAY2BGR)
-    feret_lengths, feret_widths, rectangles = [], [], []
-    contours_mask, _ = cv2.findContours(binary_mask, 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, 0, 0, 0]
-    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, 3)
-
-    fig2 = plt.figure(figsize=(6, 6))
-    plt.imshow(cv2.cvtColor(color_mask, cv2.COLOR_BGR2RGB))
-    plt.title("Feret Rectangles by Size")
-    plt.axis('off')
-    figs.append(fig2)
-
-    # Delaunay Triangulation
-    labeled_img = label(label_img)
-    props = regionprops(labeled_img)
-    centroids = np.array([p.centroid for p in props])
-    edge_lengths = []
-
-    if len(centroids) >= 3:
-        tri = Delaunay(centroids)
-        fig3 = plt.figure(figsize=(6, 6))
-        plt.imshow(label_img, cmap='gray')
-        plt.triplot(centroids[:, 1], centroids[:, 0], tri.simplices.copy(), color='red')
-        for simplex in tri.simplices:
-            for i in range(3):
-                pt1, pt2 = centroids[simplex[i]], centroids[simplex[(i + 1) % 3]]
-                dist_mm = np.linalg.norm(pt1 - pt2) * pixel_length_mm
-                edge_lengths.append(dist_mm)
-                midpoint = (pt1 + pt2) / 2
-                plt.text(midpoint[1], midpoint[0], f"{dist_mm:.1f}", color='blue', fontsize=6, ha='center')
-        plt.title("Delaunay Triangulation")
-        plt.axis('off')
-        figs.append(fig3)
-
-    # Summary Stats
-    area_mask = np.sum(binary_mask > 0)
-    area_gray = np.count_nonzero(gray_img)
-    aggregate_area_mm2 = area_mask * (pixel_length_mm ** 2)
-    total_area_mm2 = area_gray * (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
-
-    summary = f"""
-→ 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
-→ Avg Aggregate Roundness: {roundness_aggregate:.2f}
-"""
-    if edge_lengths:
-        summary += f"""
-→ Avg inter_Aggregate Distance: {np.mean(edge_lengths):.2f} mm
-→ Max inter_Aggregate Distance: {np.max(edge_lengths):.2f} mm
-"""
-
-    images = [fig_to_image(fig) for fig in figs]
-    return images[0], images[1], images[2] if len(images) > 2 else images[1], summary
-
-iface = gr.Interface(
-    fn=predict,
-    inputs=[gr.Image(label="Upload Concrete Image")],
-    outputs=[
-        gr.Image(label="Boundary and Calibration Line"),
-        gr.Image(label="Feret Rectangles by Size"),
-        gr.Image(label="Delaunay Triangulation"),
-        gr.Textbox(label="Summary Measurements")
-    ],
-    title="Concrete Aggregate Analysis App",
-    description="Upload a concrete cross-section image. The app segments aggregates, displays Feret rectangles, boundary calibration, Delaunay triangulation, and summary measurements."
-)
-
-if __name__ == "__main__":
-    iface.launch()
+import os
+import cv2
+import numpy as np
+import torch
+import matplotlib.pyplot as plt
+from scipy.spatial.distance import cdist
+from scipy.spatial import Delaunay
+from skimage.measure import label, regionprops
+import gradio as gr
+import io
+from PIL import Image
+
+# Constants
+DIA_MM = 152.4
+
+# Main processing function
+def analyze_aggregate(image_pil):
+    results = {}
+    edge_lengths = []
+
+    # Convert to OpenCV image
+    img = np.array(image_pil)
+    img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+    gray_img = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
+
+    # Simulated label (as if predicted by a model)
+    label_img = cv2.cvtColor(gray_img, cv2.COLOR_GRAY2BGR)  # Dummy label for placeholder
+    _, label_gray = cv2.threshold(gray_img, 127, 255, cv2.THRESH_BINARY)
+    binary_mask = (label_gray > 0).astype(np.uint8)
+    color_mask = cv2.cvtColor(label_gray, cv2.COLOR_GRAY2BGR)
+
+    # Pixel calibration
+    _, 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 "No contours found.", None, None, None, None
+
+    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 / DIA_MM
+    pixel_length_mm = 1 / pixels_per_mm
+    line_length_mm = pixel_diameter * pixel_length_mm
+
+    # Plot 1: Boundary and line
+    fig1, ax1 = plt.subplots(figsize=(6, 6))
+    ax1.imshow(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB))
+    ax1.plot(boundary[:, 0], boundary[:, 1], 'g', linewidth=2)
+    ax1.plot(line_pts[:, 0], line_pts[:, 1], 'r', linewidth=2)
+    ax1.set_title(f'Line Length: {line_length_mm:.2f} mm')
+    ax1.axis('off')
+
+    # Aggregate area
+    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
+
+    # Feret Rectangles
+    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, 0, 0, 0]
+    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, 3)
+
+    # Plot 2: Feret rectangles
+    fig2, ax2 = plt.subplots(figsize=(6, 6))
+    ax2.imshow(cv2.cvtColor(color_mask, cv2.COLOR_BGR2RGB))
+    ax2.set_title("Feret Rectangles by Size")
+    ax2.axis('off')
+
+    # Feret Stats
+    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
+
+    # Delaunay triangulation
+    labeled_img = label(binary_mask)
+    props = regionprops(labeled_img)
+    centroids = np.array([p.centroid for p in props])
+
+    if len(centroids) >= 3:
+        tri = Delaunay(centroids)
+        fig3, ax3 = plt.subplots(figsize=(6, 6))
+        ax3.imshow(label_gray, cmap='gray')
+        ax3.triplot(centroids[:, 1], centroids[:, 0], tri.simplices.copy(), color='red')
+
+        for simplex in tri.simplices:
+            for i in range(3):
+                pt1 = centroids[simplex[i]]
+                pt2 = centroids[(i + 1) % 3]
+                dist_px = np.linalg.norm(pt1 - pt2)
+                dist_mm = dist_px * pixel_length_mm
+                edge_lengths.append(dist_mm)
+                midpoint = (pt1 + pt2) / 2
+                ax3.text(midpoint[1], midpoint[0], f"{dist_mm:.1f}", color='blue', fontsize=6, ha='center')
+
+        ax3.set_title("Delaunay Triangulation")
+        ax3.axis('off')
+    else:
+        fig3 = plt.figure()
+        plt.text(0.5, 0.5, 'Not enough centroids for triangulation.', ha='center')
+        plt.axis('off')
+
+    # Summary text
+    summary = f"""
+→ 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
+→ Avg Aggregate Roundness: {roundness_aggregate:.2f}
+"""
+    if edge_lengths:
+        summary += f"""
+→ Avg inter-Aggregate Distance: {np.mean(edge_lengths):.2f} mm
+→ Max inter-Aggregate Distance: {np.max(edge_lengths):.2f} mm
+"""
+
+    return summary.strip(), fig1, fig2, fig3
+
+# Gradio UI
+demo = gr.Interface(
+    fn=analyze_aggregate,
+    inputs=[gr.Image(label="Upload Image")],
+    outputs=[
+        gr.Textbox(label="Summary Measurements"),
+        gr.Plot(label="Boundary and Calibration Line"),
+        gr.Plot(label="Feret Rectangles by Size"),
+        gr.Plot(label="Delaunay Triangulation")
+    ],
+    title="Aggregate Analysis from Uploaded Image",
+    description="Upload an image with circular calibration. The app will calculate size, aspect ratio, and spacing of aggregates.",
+    allow_flagging='never'
+)
+
+demo.launch()