Update app.py

app.py

@@ -1,23 +1,38 @@
-import pickle
 import cv2
 from fastai.vision.all import *
 import gradio as gr
 
-im_size = (1152, 768)
-with open("learn.pkl", 'rb') as f:
-    model = pickle.load(f)
+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)
 
-def predict_mask(file):
-    img = cv2.imdecode(np.frombuffer(file.read(), np.uint8), cv2.IMREAD_GRAYSCALE)
-    img = cv2.resize(img, im_size)
-    pred = model.predict(img)
-    return pred[0]
+dls = SegmentationDataLoaders.from_label_func(
+    ".", bs=3, fnames = fnames, label_func = label_func, codes = codes,
+    item_tfms=Resize(img_size)
+)
 
-iface = gr.Interface(fn=predict_mask, 
-                     inputs=gr.inputs.Image(label="Upload Image"), 
-                     outputs="image",
-                     capture_session=True,
-                     interpretation="default")
+learn = unet_learner(dls, resnet34)
+learn.load('learn')
 
-if __name__ == "__main__":
-    iface.launch()
+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)
+    # Convert the predicted mask to a color image
+    color_pred = pred[0].show(ctx=None, cmap='gray', alpha=None)
+    # Convert the color image to a numpy array
+    color_pred_array = np.array(color_pred)
+    # Convert the numpy array back to a PIL image
+    output_image = Image.fromarray(color_pred_array)
+    return output_image
+
+input_image = gr.inputs.Image()
+output_image = gr.outputs.Image()
+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()