finish

app.py

@@ -105,11 +105,36 @@ def sepia(input_img):
     return fig
 
 
-iface = gr.Interface(fn=sepia,
-                    inputs=gr.inputs.Image(type="pil", label="Upload an image"),
-                    outputs=gr.outputs.Image(type='pil', label="Segmentation result"),
+def compare_images(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]
+    )
+    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)
+
+    return input_img, pred_img
+
+
+demo = gr.Interface(fn=compare_images,
+                    inputs=gr.Image(shape=(1024, 1024)),
+                    outputs=["image", "image"],
                     examples=["city-1.jpg", "city-2.jpg", "city-3.jpg", "city-4.jpg", "city-5.jpg"],
                     allow_flagging='never')
 
 
-iface.launch()
+demo.launch()