Update segment.py

segment.py

@@ -1,6 +1,7 @@
 
 from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation
 from PIL import Image
+import spaces
 import torch
 from collections import defaultdict
 import matplotlib.pyplot as plt
@@ -12,6 +13,7 @@ import argparse
 import matplotlib
 import gradio as gr
 
+@spaces.GPU
 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]