Update app.py

app.py

@@ -18,18 +18,38 @@ dls = SegmentationDataLoaders.from_label_func(
 learn = unet_learner(dls, resnet34)
 learn.load('learn')
 
+def segmentImage(img_path):
+    img = cv2.imread(img_path, 0)
+    for i in range(img.shape[0]):
+        for j in range(img.shape[1]):
+            if img[i][j] > 0:
+                img[i][j] = 1
+    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)
+    labels, nlabels = ndimage.label(img)
+    colors = np.random.randint(0, 255, (nlabels + 1, 3))
+    colors[0] = 0
+    img_color = colors[labels]
+    return img_color
+
 def predict_segmentation(img):
-    # Convert the input image to grayscale
     gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
-    # Resize the image to the size of the training images
     resized_img = cv2.resize(gray_img, im_size)
-    # Predict the segmentation mask
     pred = learn.predict(resized_img)
-    scaled_pred = (pred[0].numpy() * 255).astype(np.uint8)
-    output_image = PILImage.create(scaled_pred)
-    return output_image
+    color_pred = pred[0].show(ctx=None, cmap='gray', alpha=None)
+    color_pred_array = color_pred_array.astype(np.uint8)
+    output_image = Image.fromarray(color_pred_array)
+    # Save the image to a temporary file
+    temp_file = 'temp.png'
+    output_image.save(temp_file)
+    # Call the segmentImage function
+    segmented_image = segmentImage(temp_file)
+    return output_image, segmented_image
 
 input_image = gr.inputs.Image()
-output_image = gr.outputs.Image(type='pil')
-app = gr.Interface(fn=predict_segmentation, inputs=input_image, outputs=output_image, 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')
+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.launch()