File size: 1,225 Bytes
ea326f6 1709591 ea326f6 fab2111 ea326f6 fab2111 ea326f6 fab2111 ea326f6 fab2111 d64c0c1 fab2111 cb0a970 fab2111 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
import cv2
from fastai.vision.all import *
import numpy as np
import gradio as gr
fnames = get_image_files("./albumentations/original")
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)
dls = SegmentationDataLoaders.from_label_func(
".", bs=3, fnames = fnames, label_func = label_func, codes = codes,
item_tfms=Resize(img_size)
)
learn = unet_learner(dls, resnet34)
learn.load('learn')
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] * 255).astype(np.uint8)
output_image = PILImage.create(scaled_pred)
return output_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()
|