Update app.py

app.py

@@ -71,7 +71,7 @@ def label_to_color_image(label):
         raise ValueError("label value too large.")
     return colormap[label]
 
-def draw_plot(pred_img, seg):
+def draw_plot(pred_img, seg, cursor_pos):
     fig = plt.figure(figsize=(20, 15))
 
     grid_spec = gridspec.GridSpec(1, 2, width_ratios=[6, 1])
@@ -85,14 +85,19 @@ def draw_plot(pred_img, seg):
 
     unique_labels = np.unique(seg.numpy().astype("uint8"))
     ax = plt.subplot(grid_spec[1])
-    plt.imshow(FULL_COLOR_MAP[unique_labels].astype(np.uint8), interpolation="nearest")
+    cursor_x, cursor_y = cursor_pos
+
+    mask = seg.numpy() == seg.numpy()[cursor_x, cursor_y]
+    mask_image = FULL_COLOR_MAP[mask].reshape(pred_img.shape)
+
+    plt.imshow(mask_image.astype(np.uint8), interpolation="nearest")
     ax.yaxis.tick_right()
     plt.yticks(range(len(unique_labels)), LABEL_NAMES[unique_labels])
     plt.xticks([], [])
     ax.tick_params(width=0.0, labelsize=25)
     return fig
 
-def sepia(input_img):
+def sepia(input_img, cursor_pos):
     input_img = Image.fromarray(input_img)
 
     inputs = feature_extractor(images=input_img, return_tensors="tf")
@@ -102,29 +107,17 @@ def sepia(input_img):
     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]
 
-    color_seg = np.zeros(
-        (seg.shape[0], seg.shape[1], 3), dtype=np.uint8
-    )  # height, width, 3
-    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)
-
-    fig = draw_plot(pred_img, seg)
-
-    
+    fig = draw_plot(np.array(input_img), seg, cursor_pos)
     return fig
 
 demo = gr.Interface(fn=sepia,
-                    inputs=gr.Image(shape=(400, 600)),
-                    outputs=['plot'],
-                    examples=["side-1.jpg", "side-2.jpg", "side-3.jpg", "side-4.jpg"],
+                    inputs=["image", "canvas"],
+                    outputs="plot",
+                    examples=[["side-1.jpg", [200, 300]], ["side-2.jpg", [150, 250]], ["side-3.jpg", [100, 200]], ["side-4.jpg", [250, 400]]],
+                    live=True,
                     allow_flagging='never')
 
-
 demo.launch()