File size: 2,491 Bytes
4170a88
8db4fd8
5ffae2f
5a5c67e
8db4fd8
abbed59
4170a88
 
 
 
 
 
 
 
 
5ffae2f
4170a88
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8db4fd8
 
 
 
 
 
 
 
 
 
5ffae2f
8db4fd8
 
5ffae2f
8db4fd8
 
c1ae8c8
5ffae2f
abbed59
5ffae2f
abbed59
5ffae2f
c1ae8c8
8db4fd8
 
 
 
 
 
736ffb2
 
c1ae8c8
8db4fd8
 
736ffb2
 
c1ae8c8
8db4fd8
 
 
 
 
 
5ffae2f
abbed59
 
 
8db4fd8
5ffae2f
abbed59
4170a88
5ffae2f
074b669
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
85
86
87
88
89
90
91
92
import gradio as gr
from PIL import Image
import numpy as np
import tensorflow as tf
from transformers import SegformerFeatureExtractor, TFSegformerForSemanticSegmentation
import os

feature_extractor = SegformerFeatureExtractor.from_pretrained(
    "nvidia/segformer-b1-finetuned-cityscapes-1024-1024"
)
model = TFSegformerForSemanticSegmentation.from_pretrained(
    "nvidia/segformer-b1-finetuned-cityscapes-1024-1024"
)

def ade_palette():
    """ADE20K palette that maps each class to RGB values."""
    return [
        [204, 87, 92],
        [112, 185, 212],
        [45, 189, 106],
        [234, 123, 67],
        [78, 56, 123],
        [210, 32, 89],
        [90, 180, 56],
        [155, 102, 200],
        [33, 147, 176],
        [255, 183, 76],
        [67, 123, 89],
        [190, 60, 45],
        [134, 112, 200],
        [56, 45, 189],
        [200, 56, 123],
        [87, 92, 204],
        [120, 56, 123],
        [45, 78, 123],
        [156, 200, 56]
    ]

labels_list = []

with open(r'labels.txt', 'r') as fp:
    for line in fp:
        labels_list.append(line[:-1])

colormap = np.asarray(ade_palette())

def label_to_color_image(label):
    if label.ndim != 2:
        raise ValueError("Expect 2-D input label")

    if np.max(label) >= len(colormap):
        raise ValueError("label value too large.")
    return colormap[label]

def sepia(input_text):
    # Check if the input text is a valid file path
    if not os.path.isfile(input_text):
        return "Invalid file path. Please enter a valid image file path."

    # Load the image using the input text (assumed to be a path to an image)
    input_img = Image.open(input_text)

    inputs = feature_extractor(images=input_img, return_tensors="tf")
    outputs = model(**inputs)
    logits = outputs.logits

    logits = tf.transpose(logits, [0, 2, 3, 1])
    logits = tf.image.resize(
        logits, input_img.size[::-1]
    )
    seg = tf.math.argmax(logits, axis=-1)[0]

    color_seg = np.zeros(
        (seg.shape[0], seg.shape[1], 3), dtype=np.uint8
    )
    for label, color in enumerate(colormap):
        color_seg[seg.numpy() == label, :] = color

    pred_img = np.array(input_img) * 0.5 + color_seg * 0.5
    pred_img = pred_img.astype(np.uint8)

    # Convert the image array to a Pillow (PIL) image
    pred_img = Image.fromarray(pred_img)

    return pred_img

# Define the Gradio interface
iface = gr.Interface(fn=sepia, inputs="image", outputs="image")

# Launch the Gradio app
iface.launch()