Commit
·
248a469
1
Parent(s):
9155a25
Update app.py
Browse files
app.py
CHANGED
@@ -18,18 +18,38 @@ dls = SegmentationDataLoaders.from_label_func(
|
|
18 |
learn = unet_learner(dls, resnet34)
|
19 |
learn.load('learn')
|
20 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
21 |
def predict_segmentation(img):
|
22 |
-
# Convert the input image to grayscale
|
23 |
gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
24 |
-
# Resize the image to the size of the training images
|
25 |
resized_img = cv2.resize(gray_img, im_size)
|
26 |
-
# Predict the segmentation mask
|
27 |
pred = learn.predict(resized_img)
|
28 |
-
|
29 |
-
|
30 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
31 |
|
32 |
input_image = gr.inputs.Image()
|
33 |
-
|
34 |
-
|
35 |
-
app.
|
|
|
|
18 |
learn = unet_learner(dls, resnet34)
|
19 |
learn.load('learn')
|
20 |
|
21 |
+
def segmentImage(img_path):
|
22 |
+
img = cv2.imread(img_path, 0)
|
23 |
+
for i in range(img.shape[0]):
|
24 |
+
for j in range(img.shape[1]):
|
25 |
+
if img[i][j] > 0:
|
26 |
+
img[i][j] = 1
|
27 |
+
kernel = np.ones((3,3), np.uint8)
|
28 |
+
img = cv2.erode(img, kernel, iterations=1)
|
29 |
+
img = cv2.dilate(img, kernel, iterations=1)
|
30 |
+
img = ndimage.binary_fill_holes(img).astype(int)
|
31 |
+
labels, nlabels = ndimage.label(img)
|
32 |
+
colors = np.random.randint(0, 255, (nlabels + 1, 3))
|
33 |
+
colors[0] = 0
|
34 |
+
img_color = colors[labels]
|
35 |
+
return img_color
|
36 |
+
|
37 |
def predict_segmentation(img):
|
|
|
38 |
gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
|
|
39 |
resized_img = cv2.resize(gray_img, im_size)
|
|
|
40 |
pred = learn.predict(resized_img)
|
41 |
+
color_pred = pred[0].show(ctx=None, cmap='gray', alpha=None)
|
42 |
+
color_pred_array = color_pred_array.astype(np.uint8)
|
43 |
+
output_image = Image.fromarray(color_pred_array)
|
44 |
+
# Save the image to a temporary file
|
45 |
+
temp_file = 'temp.png'
|
46 |
+
output_image.save(temp_file)
|
47 |
+
# Call the segmentImage function
|
48 |
+
segmented_image = segmentImage(temp_file)
|
49 |
+
return output_image, segmented_image
|
50 |
|
51 |
input_image = gr.inputs.Image()
|
52 |
+
output_image1 = gr.outputs.Image(type='pil')
|
53 |
+
output_image2 = gr.outputs.Image(type='pil')
|
54 |
+
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.')
|
55 |
+
app.launch()
|