Update segment.py

segment.py

@@ -11,6 +11,7 @@ import numpy as np
 import argparse
 import matplotlib
 import gradio as gr
+import cv2
 
 def load_image(image_path, left=0, right=0, top=0, bottom=0, size = 512):
     if type(image_path) is str:
@@ -52,6 +53,8 @@ def draw_panoptic_segmentation(segmentation, segments_info,save_folder=None, nos
         if torch.min(segmentation) == 0: 
             mask = segmentation==0
             mask = mask.cpu().detach().numpy()   # [512,512]   bool
+            print(mask.shape)
+            mask = cv2.resize(mask,(512,512))
             segment_label = "rest"
             color = viridis(0)
             label = f"{segment_label}-{0}"
@@ -65,6 +68,8 @@ def draw_panoptic_segmentation(segmentation, segments_info,save_folder=None, nos
             if torch.min(segmentation) != 0: 
                 segment_id -= 1
             mask = mask.cpu().detach().numpy()   # [512,512] bool
+            print(mask.shape)
+            mask = cv2.resize(mask,(512,512))
             mask_np_list.append(mask) 
             segment_label = model.config.id2label[segment['label_id']]
             instances_counter[segment['label_id']] += 1
@@ -76,6 +81,9 @@ def draw_panoptic_segmentation(segmentation, segments_info,save_folder=None, nos
             label_list.append(label)
     else:
         mask = np.full(segmentation.shape, True)
+        print(mask.shape)
+        mask = cv2.resize(mask,(512,512))
+            
         segment_label = "all"
         mask_np_list.append(mask) 
         color = viridis(0)