Add segmented image

app.py

@@ -29,15 +29,18 @@ with demo:
                 btn_next = gr.Button("Get Next Sample")
 
         with gr.Column():
-            with gr.Row():
-                info = gr.Text(label="Sample Info")
             with gr.Row():
                 img = gr.Image(label="Sample", type="numpy")
+            with gr.Row():
+                with gr.Column():
+                    info = gr.Text(label="Info")
+                with gr.Column():
+                    img_seg = gr.Image(label="Segmentation", type="numpy")
 
     btn_next.click(
         fn=sample.explorateSamples,
         inputs=[username,predictions,category,model,split,next_idx_sample], 
-        outputs=[next_idx_sample, progress, img, info])
+        outputs=[next_idx_sample, progress, img, info, img_seg])
 
 demo.queue(concurrency_count=10)
 demo.launch(debug=False)

utils.py

@@ -22,8 +22,8 @@ class SampleClass:
         img = np.array(pill_img)
         return img
 
-    def __loadPredictions(self, split, model):
 
+    def __loadPredictions(self, split, model):
         assert(split in ['test','val'])
         assert(model in ['baseline','extended'])
 
@@ -56,7 +56,6 @@ class SampleClass:
         pos_tags = sample['pos_tags'].values[0]
         plural_tks = sample['plural_tks'].values[0]
 
-
         cat_intrinsic = sample['intrinsic'].values[0]
         cat_spatial = sample['spatial'].values[0]
         cat_ordinal = sample['ordinal'].values[0]
@@ -67,10 +66,6 @@ class SampleClass:
                     ('ordinal',cat_ordinal),
                     ('relational',cat_relational),
                     ('plural',cat_plural)]
-        
-        
-        img_path = "saiapr_tc-12"+sample['file_path'].values[0].split("saiapr_tc-12")[1]
-        # print(img_path)
 
         hit = sample['hit'].values[0]
         hit_o = sample['hit_other'].values[0]
@@ -92,7 +87,11 @@ class SampleClass:
         x1_pred_o,y1_pred_o,x2_pred_o,y2_pred_o = bp_o_bbox
         # x1_pred_o,y1_pred_o,x2_pred_o,y2_pred_o = tuple(map(float,bp_o_bbox[1:-1].split(",")))
 
-        # Plot
+        # Create Fig with predictions
+        img_path = "saiapr_tc-12"+sample['file_path'].values[0].split("saiapr_tc-12")[1]
+        img_seg_path = img_path.replace("images","segmented_images")
+
+
         fig, ax = plt.subplots(1)
         ax.imshow(self.__get(img_path), interpolation='bilinear')
 
@@ -128,7 +127,9 @@ class SampleClass:
         data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
         w, h = fig.canvas.get_width_height()
         img = data.reshape((int(h), int(w), -1))
-        return info, img
+
+
+        return info, img, self.__get(img_seg_path)
 
     def explorateSamples(self, 
             username,
@@ -160,8 +161,8 @@ class SampleClass:
 
         next_idx_sample = user_ids[username][ min(id_+1, len(user_ids[username])-1) ]
         progress = {f"{id_}/{len(user_ids[username])-1}":id_/(len(user_ids[username])-1)}
-        info, img = self.__getSample(user_ids[username][id_])
+        info, img, img_seg = self.__getSample(user_ids[username][id_])
         info = "".join([str(k)+":\t"+str(v)+"\n" for k,v in list(info.items())[1:]]).strip()
         
-        return (gr.Number.update(value=next_idx_sample),progress,img,info)
+        return (gr.Number.update(value=next_idx_sample),progress,img,info,img_seg)