Update segment.py

segment.py

@@ -1,9 +1,9 @@
 
 from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
 from PIL import Image
-import spaces
 import torch
 from collections import defaultdict
+from functools import lru_cache
 import matplotlib.pyplot as plt
 from matplotlib import cm
 import matplotlib.patches as mpatches
@@ -14,6 +14,15 @@ import matplotlib
 import gradio as gr
 
 
+@lru_cache(maxsize=1)
+def get_segmentation_model():
+    model_id = "facebook/mask2former-swin-base-coco-panoptic"
+    processor = AutoImageProcessor.from_pretrained(model_id)
+    model = Mask2FormerForUniversalSegmentation.from_pretrained(model_id)
+    model.eval()
+    return processor, model
+
+
 def load_image(image_path, left=0, right=0, top=0, bottom=0, size = 512):
     if type(image_path) is str:
         image = np.array(Image.open(image_path))[:, :, :3]
@@ -94,17 +103,16 @@ def draw_panoptic_segmentation(segmentation, segments_info,save_folder=None, nos
     # plt.savefig(os.path.join(save_folder, 'mask_clear.png'), dpi=500)
     ax.legend(handles=handles)
     plt.savefig(os.path.join(save_folder, 'seg_init.png'), dpi=500 )
+    plt.close(fig)
     print("; ".join(label_list))
     return mask_np_list,label_list
 
 
-@spaces.GPU(duration=10)
 def run_segmentation(image, name="example_tmp", size = 512, noseg=False):
 
     base_folder_path = "."
 
-    processor = AutoImageProcessor.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")
-    model = Mask2FormerForUniversalSegmentation.from_pretrained("facebook/mask2former-swin-base-coco-panoptic")
+    processor, model = get_segmentation_model()
 
 
     # input_folder = os.path.join(base_folder_path, name )
@@ -117,7 +125,7 @@ def run_segmentation(image, name="example_tmp", size = 512, noseg=False):
     os.makedirs(name, exist_ok=True)
     #image.save(os.path.join(name,"img_{}.png".format(size)))
     inputs = processor(image, return_tensors="pt")
-    with torch.no_grad():
+    with torch.inference_mode():
         outputs = model(**inputs)
 
     panoptic_segmentation = processor.post_process_panoptic_segmentation(outputs, target_sizes=[image.size[::-1]])[0]