a

app.py

@@ -3,9 +3,10 @@ import gradio as gr
 from transformers import SegformerFeatureExtractor, SegformerForSemanticSegmentation
 import matplotlib.pyplot as plt
 from matplotlib import gridspec
-from PIL import Image
 import numpy as np
 import tensorflow as tf
+from PIL import Image
+from io import BytesIO
 import requests
 
 #
@@ -27,60 +28,58 @@ model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-fi
 # outputs = model(**inputs)
 # logits = outputs.logits  # shape (batch_size, num_labels, height/4, width/4)
 
-def my_palette():
-    return [
-        [131, 162, 255],
-        [180, 189, 255],
-        [255, 227, 187],
-        [255, 210, 143],
-        [248, 117, 170],
-        [255, 223, 223],
-        [255, 246, 246],
-        [174, 222, 252],
-        [150, 194, 145],
-        [255, 219, 170],
-        [244, 238, 238],
-        [50, 38, 83],
-        [128, 98, 214],
-        [146, 136, 248],
-        [255, 210, 215],
-        [255, 152, 152],
-        [162, 103, 138],
-        [63, 29, 56]
-    ]
-
-
 labels_list = []
-
-with open(r"labels.txt", "r") as fp:
+with open("labels.txt", "r") as fp:
     for line in fp:
         labels_list.append(line[:-1])
 
-colormap = np.asarray(my_palette())
+colormap = np.asarray([
+    [131, 162, 255],
+    [180, 189, 255],
+    [255, 227, 187],
+    [255, 210, 143],
+    [248, 117, 170],
+    [255, 223, 223],
+    [255, 246, 246],
+    [174, 222, 252],
+    [150, 194, 145],
+    [255, 219, 170],
+    [244, 238, 238],
+    [50, 38, 83],
+    [128, 98, 214],
+    [146, 136, 248],
+    [255, 210, 215],
+    [255, 152, 152],
+    [162, 103, 138],
+    [63, 29, 56]
+])
+
+# with open(r"labels.txt", "r") as fp:
+#     for line in fp:
+#         labels_list.append(line[:-1])
+#
+# colormap = np.asarray(my_palette())
 
 
 def greet(input_img):
+    input_img = Image.open(BytesIO(input_img))
+
     inputs = feature_extractor(images=input_img, return_tensors="pt")
     outputs = model(**inputs)
     logits = outputs.logits
 
     logits = tf.transpose(logits, [0, 2, 3, 1])
+    logits = tf.image.resize(logits, input_img.size[::-1])
 
-    logits = tf.image.resize(
-        logits, input_img.size[::-1]
-    )  # We reverse the shape of `image` because `image.size` returns width and height.
     seg = tf.math.argmax(logits, axis=-1)[0]
-
-    color_seg = np.zeros(
-        (seg.shape[0], seg.shape[1], 3), dtype=np.uint8
-    )  # height, width, 3
+    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
 
-    # Show image + mask
     pred_img = np.array(input_img) * 0.5 + color_seg * 0.5
     pred_img = pred_img.astype(np.uint8)
 
+    # Draw the plot
     fig = draw_plot(pred_img, seg)
     return fig
 
@@ -93,6 +92,7 @@ def draw_plot(pred_img, seg):
     plt.subplot(grid_spec[0])
     plt.imshow(pred_img)
     plt.axis("off")
+
     LABEL_NAMES = np.asarray(labels_list)
     FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
     FULL_COLOR_MAP = label_to_color_image(FULL_LABEL_MAP)
@@ -121,5 +121,6 @@ iface = gr.Interface(
     inputs="image",
     outputs=["plot"],
     examples=["image (1).jpg", "image (2).jpg", "image (3).jpg", "image (4).jpg", "image (5).jpg"],
-    allow_flagging="never")
+    allow_flagging="never"
+)
 iface.launch(share=True)

requirements.txt

@@ -1,6 +1,7 @@
-torch
-transformers
-tensorflow
-numpy
-Image
-matplotlib
+gradio==2.3.0
+torch==1.10.0
+transformers==4.13.0
+matplotlib==3.4.3
+numpy==1.21.2
+pillow==8.3.2
+tensorflow==2.6.0