Update app.py

app.py

@@ -1,68 +1,120 @@
 import os
 import cv2
-import numpy as np
 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 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)
+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 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)
 
-    # 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
+        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 / DIA_MM
+    pixels_per_mm = pixel_diameter / DIAMETER_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')
+    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)
 
-    # 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
+    binary_mask = (label_img > 127).astype(np.uint8)
+    color_mask = cv2.cvtColor(label_img, cv2.COLOR_GRAY2BGR)
 
-    # Feret Rectangles
-    feret_lengths, feret_widths = [], []
-    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:
@@ -74,9 +126,8 @@ def analyze_aggregate(image_pil):
             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)]
-
+    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]
@@ -85,80 +136,78 @@ def analyze_aggregate(image_pil):
         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
+    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)
 
-    # Delaunay triangulation
     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)
-        fig3, ax3 = plt.subplots(figsize=(6, 6))
-        ax3.imshow(label_gray, cmap='gray')
-        ax3.triplot(centroids[:, 1], centroids[:, 0], tri.simplices.copy(), color='red')
-
+        plt.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]
+                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)
                 midpoint = (pt1 + pt2) / 2
-                ax3.text(midpoint[1], midpoint[0], f"{dist_mm:.1f}", color='blue', fontsize=6, ha='center')
+                plt.text(midpoint[1], midpoint[0], f"{dist_mm:.1f}", color="blue", fontsize=6, ha="center")
+        plt.title("Delaunay Triangulation")
+    else:
+        plt.title("Not Enough Aggregates for Triangulation")
+    plt.axis("off")
+    img3 = fig_to_image(fig3)
 
-        ax3.set_title("Delaunay Triangulation")
-        ax3.axis('off')
+    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:
-        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}
+        avg_feret_length_mm = avg_feret_width_mm = max_feret_length_mm = roundness_aggregate = 0
+
+    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
-→ Max inter-Aggregate Distance: {np.max(edge_lengths):.2f} mm
-"""
+        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 summary.strip(), fig1, fig2, fig3
+    return img1, img2, img3, summary
 
-# Gradio UI
 demo = gr.Interface(
-    fn=analyze_aggregate,
-    inputs=[gr.Image(label="Upload Image")],
+    fn=analyze,
+    inputs=gr.Image(type="pil", label="Upload Concrete 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")
+        gr.Image(label="Boundary & Calibration Line"),
+        gr.Image(label="Feret Rectangles"),
+        gr.Image(label="Delaunay Triangulation"),
+        gr.Markdown(label="Summary Measurements"),
     ],
-    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'
+    title="Concrete Aggregate Analysis App",
+    description="Upload a concrete image. The model will segment aggregates and analyze their distribution and shape.",
 )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()