Update app.py

app.py

@@ -1,18 +1,12 @@
 import gradio as gr
 
+import numpy as np
+import cv2
 from matplotlib import gridspec
 import matplotlib.pyplot as plt
-import numpy as np
-from PIL import Image
-import tensorflow as tf
-from transformers import SegformerFeatureExtractor, TFSegformerForSemanticSegmentation
+import onnxruntime as ort
 
-feature_extractor = SegformerFeatureExtractor.from_pretrained(
-    "nvidia/segformer-b5-finetuned-ade-640-640"
-)
-model = TFSegformerForSemanticSegmentation.from_pretrained(
-    "nvidia/segformer-b5-finetuned-ade-640-640"
-)
+import wget
 
 def ade_palette():
     """ADE20K palette that maps each class to RGB values."""
@@ -169,14 +163,17 @@ def ade_palette():
         [92, 0, 255],
     ]
 
+url='https://github.com/deep-diver/segformer-tf-transformers/releases/download/1.0/segformer-b5-finetuned-ade-640-640.onnx'
 labels_list = []
+colormap = np.asarray(ade_palette())
+
+model_path = wget.download(url)
+sess = ort.InferenceSession(model_path)
 
 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")
@@ -199,7 +196,7 @@ def draw_plot(pred_img, seg):
     FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
     FULL_COLOR_MAP = label_to_color_image(FULL_LABEL_MAP)
 
-    unique_labels = np.unique(seg.numpy().astype("uint8"))
+    unique_labels = np.unique(seg)
     ax = plt.subplot(grid_spec[1])
     plt.imshow(FULL_COLOR_MAP[unique_labels].astype(np.uint8), interpolation="nearest")
     ax.yaxis.tick_right()
@@ -209,18 +206,17 @@ def draw_plot(pred_img, seg):
     return fig
 
 def sepia(input_img):
-    input_img = Image.fromarray(input_img)
-
-    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]
-    )  # We reverse the shape of `image` because `image.size` returns width and height.
-    seg = tf.math.argmax(logits, axis=-1)[0]
-
+    img = cv2.imread(input_img).astype(np.float32)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    img_batch = np.expand_dims(img, axis=0)
+    img_batch = np.transpose(img_batch, (0, 3, 1, 2))
+    
+    logits = sess.run(None, {"pixel_values": img_batch})[0]
+
+    logits = np.transpose(logits, (0, 2, 3, 1))
+    seg = np.argmax(logits, axis=-1)[0].astype('float32')
+    seg = cv2.resize(seg, (640, 640)).astype('uint8')
+    
     color_seg = np.zeros(
         (seg.shape[0], seg.shape[1], 3), dtype=np.uint8
     )  # height, width, 3
@@ -232,16 +228,15 @@ def sepia(input_img):
     color_seg = color_seg[..., ::-1]
 
     # Show image + mask
-    pred_img = np.array(input_img) * 0.5 + color_seg * 0.5
+    pred_img = img * 0.5 + color_seg * 0.5
     pred_img = pred_img.astype(np.uint8)    
 
     fig = draw_plot(pred_img, seg)
     return fig
 
 demo = gr.Interface(sepia, 
-                    gr.Image(shape=(200, 200)), 
+                    gr.inputs.Image(type="filepath", shape=(640, 640)), 
                     outputs=['plot'], 
-#                    examples=["ADE_val_00000001.jpeg"],
                     allow_flagging='never')
 
 demo.launch()