Update app.py

app.py

@@ -5,19 +5,29 @@ import gradio as gr
 from scipy import ndimage
 
 fnames = get_image_files("./albumentations/original")
-def label_func(fn): return "./albumentations/labelled/"f"{fn.stem}.png"
-codes = np.loadtxt('labels.txt', dtype=str)
+
+
+def label_func(fn):
+    return "./albumentations/labelled/" f"{fn.stem}.png"
+
+
+codes = np.loadtxt("labels.txt", dtype=str)
 w, h = 768, 1152
-img_size = (w,h)
-im_size = (h,w)
+img_size = (w, h)
+im_size = (h, w)
 
 dls = SegmentationDataLoaders.from_label_func(
-    ".", bs=3, fnames = fnames, label_func = label_func, codes = codes,
-    item_tfms=Resize(img_size)
+    ".",
+    bs=3,
+    fnames=fnames,
+    label_func=label_func,
+    codes=codes,
+    item_tfms=Resize(img_size),
 )
 
 learn = unet_learner(dls, resnet34)
-learn.load('learn')
+learn.load("learn")
+
 
 def segmentImage(img_path):
     img = cv2.imread(img_path, 0)
@@ -25,7 +35,7 @@ def segmentImage(img_path):
         for j in range(img.shape[1]):
             if img[i][j] > 0:
                 img[i][j] = 1
-    kernel = np.ones((3,3), np.uint8)
+    kernel = np.ones((3, 3), np.uint8)
     img = cv2.erode(img, kernel, iterations=1)
     img = cv2.dilate(img, kernel, iterations=1)
     img = ndimage.binary_fill_holes(img).astype(int)
@@ -56,10 +66,21 @@ def segmentImage(img_path):
         for j in range(img.shape[1]):
             if labels[i][j] != 0:
                 gradient_img[i][j] = colors[labels[i][j]]
+    Sum = 0
+    count = 0
+    for i in range(len(new_sizes)):
+        if new_sizes[i] > 250 * c * c:
+            Sum += new_sizes[i]
+            count += 1
     colors = np.random.randint(0, 255, (nlabels + 1, 3))
     colors[0] = 0
     img_color = colors[labels]
-    return img_color, gradient_img
+    return (
+        img_color,
+        gradient_img,
+        "Average Area of grains: " + str(Sum / count) + " µm^2",
+    )
+
 
 def predict_segmentation(img):
     gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
@@ -68,14 +89,23 @@ def predict_segmentation(img):
     scaled_pred = (pred[0].numpy() * 255).astype(np.uint8)
     output_image = PILImage.create(scaled_pred)
     # Save the image to a temporary file
-    temp_file = 'temp.png'
+    temp_file = "temp.png"
     output_image.save(temp_file)
     # Call the segmentImage function
-    segmented_image, gradient_image = segmentImage(temp_file)
-    return output_image, segmented_image, gradient_image
+    segmented_image, gradient_image, avg_area = segmentImage(temp_file)
+    return output_image, segmented_image, gradient_image, avg_area
+
 
 input_image = gr.inputs.Image()
-output_image1 = gr.outputs.Image(type='pil')
-output_image2 = gr.outputs.Image(type='pil')
-app = gr.Interface(fn=predict_segmentation, inputs=input_image, outputs=[output_image1, output_image2, output_image3], title='Microstructure Segmentation', description='Segment the input image into grain and background.')
-app.launch()
+output_image1 = gr.outputs.Image(type="pil")
+output_image2 = gr.outputs.Image(type="pil")
+output_image3 = gr.outputs.Image(type="pil")
+output_image4 = gr.outputs.Textbox()
+app = gr.Interface(
+    fn=predict_segmentation,
+    inputs=input_image,
+    outputs=[output_image1, output_image2, output_image3, output_image4],
+    title="Microstructure Segmentation",
+    description="Segment the input image into grain and background.",
+)
+app.launch()