Update app.py

app.py

@@ -30,10 +30,36 @@ def segmentImage(img_path):
     img = cv2.dilate(img, kernel, iterations=1)
     img = ndimage.binary_fill_holes(img).astype(int)
     labels, nlabels = ndimage.label(img)
+
+    # Get grain sizes
+    sizes = ndimage.sum(img, labels, range(nlabels + 1))
+    scale_factor = 3072 / 1152
+    c = 0.4228320313
+    # Divide sizes by pixel_to_micrometer to get the sizes in micrometers and store them in a list new_sizes
+    new_sizes = [size * scale_factor * scale_factor * c * c for size in sizes]
+    # Round the grain sizes to 2 decimal places
+    new_sizes = [round(size, 2) for size in new_sizes]
+    gradient_img = np.zeros((img.shape[0], img.shape[1], 3), np.uint8)
+    colors = []
+    for i in range(len(new_sizes)):
+        if new_sizes[i] < 250 * c * c:
+            colors.append((255, 255, 255))
+        elif new_sizes[i] < 7500 * c * c:
+            colors.append((2, 106, 248))
+        elif new_sizes[i] < 20000 * c * c:
+            colors.append((0, 255, 107))
+        elif new_sizes[i] < 45000 * c * c:
+            colors.append((255, 201, 60))
+        else:
+            colors.append((255, 0, 0))
+    for i in range(img.shape[0]):
+        for j in range(img.shape[1]):
+            if labels[i][j] != 0:
+                gradient_img[i][j] = colors[labels[i][j]]
     colors = np.random.randint(0, 255, (nlabels + 1, 3))
     colors[0] = 0
     img_color = colors[labels]
-    return img_color
+    return img_color, gradient_img
 
 def predict_segmentation(img):
     gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
@@ -45,11 +71,11 @@ def predict_segmentation(img):
     temp_file = 'temp.png'
     output_image.save(temp_file)
     # Call the segmentImage function
-    segmented_image = segmentImage(temp_file)
-    return output_image, segmented_image
+    segmented_image, gradient_image = segmentImage(temp_file)
+    return output_image, segmented_image, gradient_image
 
 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], title='Microstructure Segmentation', description='Segment the input image into grain and background.')
+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()