dhanushreddy29 commited on
Commit
248a469
·
1 Parent(s): 9155a25

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +29 -9
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
- scaled_pred = (pred[0].numpy() * 255).astype(np.uint8)
29
- output_image = PILImage.create(scaled_pred)
30
- return output_image
 
 
 
 
 
 
31
 
32
  input_image = gr.inputs.Image()
33
- output_image = gr.outputs.Image(type='pil')
34
- app = gr.Interface(fn=predict_segmentation, inputs=input_image, outputs=output_image, title='Microstructure Segmentation', description='Segment the input image into grain and background.')
35
- app.launch()
 
 
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()