Matiullah00999/Aggregate_Analysis

1import os2import cv23import torch4import numpy as np5import gradio as gr6import matplotlib.pyplot as plt7import pandas as pd8from glob import glob9from PIL import Image10from skimage.measure import regionprops, label11from scipy.spatial.distance import cdist12from scipy.spatial import Delaunay13from io import BytesIO14from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas15import segmentation_models_pytorch as smp16 17# Configuration18DEVICE = "cuda" if torch.cuda.is_available() else "cpu"19DIAMETER_MM = 152.420MIN_SIZE = 25621 22class PetModel(torch.nn.Module):23    def __init__(self, arch, encoder_name, in_channels, out_classes, **kwargs):24        super().__init__()25        self.model = smp.create_model(26            arch, encoder_name, in_channels=in_channels, classes=out_classes, **kwargs27        )28        params = smp.encoders.get_preprocessing_params(encoder_name)29        self.register_buffer("std", torch.tensor(params["std"]).view(1, 3, 1, 1))30        self.register_buffer("mean", torch.tensor(params["mean"]).view(1, 3, 1, 1))31 32    def forward(self, image):33        image = (image - self.mean) / self.std34        return self.model(image)35 36def preprocess_image(image, min_size=MIN_SIZE):37    image = np.array(image)38    if len(image.shape) == 2:39        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)40    elif image.shape[2] == 4:41        image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)42    elif image.shape[2] == 1:43        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)44 45    original_size = image.shape[:2]46    h, w = image.shape[:2]47    if h < min_size or w < min_size:48        new_size = (max(w, min_size), max(h, min_size))49        image = cv2.resize(image, new_size, interpolation=cv2.INTER_LINEAR)50 51    image = image.astype(np.float32) / 255.052    image = torch.tensor(image).permute(2, 0, 1).unsqueeze(0)53    return image, original_size54 55def postprocess_output(output, original_size):56    prob_mask = output.sigmoid()57    pred_mask = (prob_mask > 0.5).float()58    pred_mask = pred_mask.squeeze().cpu().numpy()59    if pred_mask.shape != original_size:60        pred_mask = cv2.resize(pred_mask, (original_size[1], original_size[0]), interpolation=cv2.INTER_NEAREST)61    return pred_mask62 63def load_model(model_path):64    model = PetModel("unet", "efficientnet-b5", in_channels=3, out_classes=1)65    model.load_state_dict(torch.load(model_path, map_location=DEVICE))66    model = model.to(DEVICE)67    model.eval()68    return model69 70model = load_model("segmentation_model_final.pth")71 72csv_output_path = "measurement_summary.csv"73 74def fig_to_image(fig):75    buf = BytesIO()76    canvas = FigureCanvas(fig)77    canvas.print_png(buf)78    buf.seek(0)79    return Image.open(buf)80 81def analyze(image):82    input_tensor, original_size = preprocess_image(image)83    input_tensor = input_tensor.to(DEVICE)84 85    with torch.no_grad():86        output = model(input_tensor)87 88    prediction_mask = postprocess_output(output, original_size)89    image_np = np.array(image)90    gray_img = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)91    label_img = (prediction_mask * 255).astype(np.uint8)92 93    _, bw = cv2.threshold(gray_img, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)94    contours, _ = cv2.findContours(bw, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)95    contours = sorted(contours, key=cv2.contourArea, reverse=True)96    if not contours:97        return None, None, None, "No contour found."98 99    boundary = contours[0].squeeze()100    dist_matrix = cdist(boundary, boundary)101    i, j = np.unravel_index(np.argmax(dist_matrix), dist_matrix.shape)102    line_pts = np.array([boundary[i], boundary[j]])103    pixel_diameter = np.linalg.norm(boundary[i] - boundary[j])104    pixels_per_mm = pixel_diameter / DIAMETER_MM105    pixel_length_mm = 1 / pixels_per_mm106    line_length_mm = pixel_diameter * pixel_length_mm107 108    fig1 = plt.figure(figsize=(6, 6))109    plt.imshow(image_np)110    plt.plot(boundary[:, 0], boundary[:, 1], 'g', linewidth=2)111    plt.plot(line_pts[:, 0], line_pts[:, 1], 'r', linewidth=2)112    plt.title(f"Calibration Line: {line_length_mm:.2f} mm")113    plt.axis("off")114    img1 = fig_to_image(fig1)115 116    binary_mask = (label_img > 127).astype(np.uint8)117    color_mask = cv2.cvtColor(label_img, cv2.COLOR_GRAY2BGR)118 119    feret_lengths, feret_widths, rectangles = [], [], []120    contours_mask, _ = cv2.findContours(binary_mask * 255, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)121    for cnt in contours_mask:122        if len(cnt) >= 5:123            rect = cv2.minAreaRect(cnt)124            box = cv2.boxPoints(rect).astype(np.intp)125            width, height = rect[1]126            feret_length = max(width, height)127            feret_lengths.append(feret_length)128            feret_widths.append(min(width, height))129            rectangles.append((box, feret_length))130 131    thresholds = np.percentile(feret_lengths, [20, 40, 60, 80]) if feret_lengths else [0]*4132    colors = [(0,0,255),(0,128,255),(0,255,255),(0,255,0),(255,0,0)]133    for box, length in rectangles:134        if length <= thresholds[0]: color = colors[0]135        elif length <= thresholds[1]: color = colors[1]136        elif length <= thresholds[2]: color = colors[2]137        elif length <= thresholds[3]: color = colors[3]138        else: color = colors[4]139        cv2.drawContours(color_mask, [box], 0, color, 8)140 141    fig2 = plt.figure(figsize=(6, 6))142    plt.imshow(cv2.cvtColor(color_mask, cv2.COLOR_BGR2RGB))143    plt.title("Feret Rectangles (Colored by Size)")144    plt.axis("off")145    img2 = fig_to_image(fig2)146 147    labeled_img = label(binary_mask)148    props = regionprops(labeled_img)149    centroids = np.array([p.centroid for p in props])150 151    edge_lengths = []152    fig3 = plt.figure(figsize=(6, 6))153    plt.imshow(label_img, cmap="gray")154    if len(centroids) >= 3:155        tri = Delaunay(centroids)156        plt.triplot(centroids[:, 1], centroids[:, 0], tri.simplices.copy(), color="red", linewidth=1)157        for simplex in tri.simplices:158            for i in range(3):159                pt1 = centroids[simplex[i]]160                pt2 = centroids[simplex[(i + 1) % 3]]161                dist_px = np.linalg.norm(pt1 - pt2)162                dist_mm = dist_px * pixel_length_mm163                edge_lengths.append(dist_mm)164        plt.title("Delaunay Triangulation")165    else:166        plt.title("Not Enough Aggregates for Triangulation")167    plt.axis("off")168    img3 = fig_to_image(fig3)169 170    num_white_pixels = np.sum(binary_mask == 1)171    num_nonblack_pixels = np.count_nonzero(gray_img)172    aggregate_area_mm2 = num_white_pixels * (pixel_length_mm ** 2)173    total_area_mm2 = num_nonblack_pixels * (pixel_length_mm ** 2)174    aggregate_ratio = aggregate_area_mm2 / total_area_mm2 if total_area_mm2 > 0 else 0175 176    if feret_lengths:177        avg_feret_length_mm = np.mean(feret_lengths) * pixel_length_mm178        avg_feret_width_mm = np.mean(feret_widths) * pixel_length_mm179        max_feret_length_mm = np.max(feret_lengths) * pixel_length_mm180        roundness_aggregate = avg_feret_length_mm / avg_feret_width_mm181    else:182        avg_feret_length_mm = avg_feret_width_mm = max_feret_length_mm = roundness_aggregate = 0183 184    # Save to CSV185    data = {186        "Pixel_Size_mm_per_pixel": [pixel_length_mm],187        "Aggregate_Area_mm2": [aggregate_area_mm2],188        "Aggregate_Ratio": [aggregate_ratio],189        "Avg_Length_mm": [avg_feret_length_mm],190        "Avg_Width_mm": [avg_feret_width_mm],191        "Max_Length_mm": [max_feret_length_mm],192        "Roundness": [roundness_aggregate],193        "Avg_Dist_mm": [np.mean(edge_lengths) if edge_lengths else 0],194        "Max_Dist_mm": [np.max(edge_lengths) if edge_lengths else 0]195    }196    df = pd.DataFrame(data)197    df.to_csv(csv_output_path, index=False)198 199    summary = f"""📏 **Measurements Summary**:200- Pixel Size: `{pixel_length_mm:.4f}` mm/pixel201- Aggregate Area: `{aggregate_area_mm2:.2f}` mm²202- Aggregate Ratio: `{aggregate_ratio:.4f}`203- Avg Aggregate Length: `{avg_feret_length_mm:.2f}` mm204- Avg Aggregate Width: `{avg_feret_width_mm:.2f}` mm205- Max Aggregate Length: `{max_feret_length_mm:.2f}` mm206- Aggregate Roundness: `{roundness_aggregate:.2f}`207"""208    if edge_lengths:209        summary += f"- Avg Inter-Aggregate Distance: `{np.mean(edge_lengths):.2f}` mm\n"210        summary += f"- Max Inter-Aggregate Distance: `{np.max(edge_lengths):.2f}` mm\n"211 212    return img1, img2, img3, summary213 214demo = gr.Interface(215    fn=analyze,216    inputs=gr.Image(type="pil", label="Upload Concrete Image"),217    outputs=[218        gr.Image(label="Boundary & Calibration Line"),219        gr.Image(label="Feret Rectangles"),220        gr.Image(label="Delaunay Triangulation"),221        gr.Textbox(label="Measurements Summary")222    ],223    title="Concrete Aggregate Analysis App",224    description="Upload a concrete image. The model will segment aggregates and analyze their distribution and shape."225)226 227if __name__ == "__main__":228    demo.launch()229