File size: 1,749 Bytes
4ae8579
 
 
0e8f4b9
4ae8579
8fc5d37
4ae8579
8fc5d37
4ae8579
 
 
 
 
d728d1b
8fc5d37
 
d728d1b
0e8f4b9
8fc5d37
 
 
d728d1b
 
8fc5d37
 
 
 
 
 
d728d1b
8fc5d37
 
d728d1b
8fc5d37
 
 
 
d728d1b
e8ef14f
 
8fc5d37
4ae8579
 
 
d728d1b
0e8f4b9
 
 
2b8f512
d728d1b
0e8f4b9
 
 
 
 
 
 
 
 
2b8f512
0e8f4b9
4ae8579
d728d1b
4ae8579
 
0e8f4b9
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
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):
    image = PILImage.create(image)
    prediction, _, _ = LEARNER.predict(image)

    print("Prediction shape:", prediction.shape)
    print("Unique values:", numpy.unique(prediction))

    prediction_array = numpy.array(prediction, dtype=numpy.uint8)

    colormap = _cm.get_cmap('jet', numpy.max(
        prediction_array) + 1)
    colored_mask = colormap(prediction_array)[:, :, :3]

    image_array = numpy.array(image).astype(numpy.float32) / 255.0

    overlay = (image_array * 0.7) + (colored_mask * 0.3)  # Adjust transparency

    overlay = (overlay * 255).astype(numpy.uint8)

    return overlay


demo = gradio.Interface(
    segment_image,
    inumpyuts=gradio.Image(type='pil'),
    outputs=gradio.Image(type='numpy'),
    examples=[str(image) for image in TEST_IMAGES_PATH.iterdir()]
)
demo.launch()