I am trying to convert the Image Segmentation example from lesson 1 into a Gradio app and the generated image looks nothing like I had in my notebook.
Has anyone converted the results of predict into an image for Gradio to display in the output?
This is the function I am using
from fastai.vision.all import *
import gradio as gr
import timm
from label_helper import LabelHelper
learn = load_learner(‘model.pkl’)
def segment_image(img):
if isinstance(img, np.ndarray):
img = Image.fromarray(img)
# Convert the input image to RGB mode
img = img.convert("RGB")
# Create PIL image from the input and get the segmentation results
img = PILImage.create(img)
results = learn.predict(img)
# Get the mask and normalize it
mask = results[0].numpy()
mask = (mask * 255).astype(np.uint8)
# Convert mask to RGB and ensure it matches the size of the original image
mask_img = Image.fromarray(mask).convert("RGB").resize(img.size, Image.NEAREST)
# Ensure the mask image has the same size as the original image
if mask_img.size != img.size:
mask_img = mask_img.resize(img.size, Image.NEAREST)
# Overlay the mask on the original image
overlayed_img = Image.blend(img, mask_img, alpha=0.7)
return overlayed_img
