Spaces:
Sleeping
Sleeping
File size: 2,450 Bytes
4ae8579 fa09dc1 4ae8579 8fc5d37 4ae8579 8fc5d37 4ae8579 d728d1b 8fc5d37 d728d1b 0e8f4b9 8fc5d37 d728d1b 8fc5d37 d728d1b 8fc5d37 d728d1b 8fc5d37 d728d1b e8ef14f 8fc5d37 4ae8579 6d3ae22 0342b70 6d3ae22 0342b70 6d3ae22 8103efe 6d3ae22 0e8f4b9 6d3ae22 0e8f4b9 4ae8579 d728d1b 4ae8579 990b65c 4ae8579 |
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 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
import gradio
import numpy
from matplotlib import cm
from pathlib import Path
from PIL import Image
from fastai.vision.all import load_learner, PILImage, PILMask
MODEL_PATH = Path('.') / 'models'
TEST_IMAGES_PATH = Path('.') / 'test'
def preprocess_mask(file_name):
"""Ensures masks are in grayscale format and removes transparency."""
mask_path = Path(
'/kaggle/inumpyut/car-segmentation/car-segmentation/masks') / file_name.name
mask = Image.open(mask_path)
if mask.mode == 'P':
mask = mask.convert('RGBA')
if mask.mode != 'RGBA':
mask = mask.convert('RGBA')
mask_data = mask.getdata()
new_mask_data = [
(r, g, b, 255) if a > 0 else (0, 0, 0, 255)
for r, g, b, a in mask_data
]
mask.putdata(new_mask_data)
return PILMask.create(mask.convert('L'))
LEARNER = load_learner(MODEL_PATH / 'car-segmentation_v1.pkl')
def segment_image(image):
# Store original size
original_size = image.size # (width, height)
# Resize the input image to 400x400 for the model
resized_image = image.resize((400, 400))
resized_image = PILImage.create(resized_image)
# Get the prediction from the model
prediction, _, _ = LEARNER.predict(resized_image)
# Convert prediction to a NumPy array
prediction_array = numpy.asarray(prediction, dtype=numpy.uint8)
# Resize the mask back to the original image size
prediction_resized = Image.fromarray(prediction_array).resize(original_size, Image.NEAREST)
prediction_resized = numpy.array(prediction_resized)
# Apply a colormap for visualization (using the public API)
colormap = cm._colormaps['jet']
# Normalize the mask and apply the colormap (result is in float [0,1])
colored_mask = colormap(prediction_resized / numpy.max(prediction_resized))[:, :, :3]
# Convert the original image to a NumPy array and normalize to [0,1]
image_array = numpy.array(image).astype(numpy.float32) / 255.0
# Blend the original image and the colored mask
overlay = (image_array * 0.7) + (colored_mask * 0.3)
# Convert the blended image back to uint8
overlay = (overlay * 255).astype(numpy.uint8)
return overlay
demo = gradio.Interface(
segment_image,
inputs=gradio.Image(type='pil'),
outputs=gradio.Image(type='numpy'),
examples=[str(image) for image in TEST_IMAGES_PATH.iterdir()]
)
demo.launch()
|