Update app.py

app.py

@@ -29,6 +29,7 @@ from PIL import Image
 import torch
 import requests
 import tensorflow as tf
+import spaces
 
 
 def _normalized_to_pixel_coordinates(
@@ -50,7 +51,6 @@ def _normalized_to_pixel_coordinates(
   return x_px, y_px
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
 pipe = StableDiffusionInpaintPipeline.from_pretrained(
     "stabilityai/stable-diffusion-2-inpainting",
     torch_dtype=torch.float16,
@@ -68,34 +68,8 @@ base_options = python.BaseOptions(model_asset_path='model.tflite')
 options = vision.ImageSegmenterOptions(base_options=base_options,
                                        output_category_mask=True)
 
-def get_bounding_box(mask):
-    """Generate bounding box coordinates from a binary mask."""
-    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-    if contours:
-        x, y, w, h = cv2.boundingRect(contours[0])
-        return x, y, x + w, y + h
-    return 0, 0, mask.shape[1], mask.shape[0]
-
-
-def example_segmentation_function(image_file_path, x, y):
-    OVERLAY_COLOR = (255, 105, 180)  # Rose
-    base_options = python.BaseOptions(model_asset_path='model.tflite')
-    options = vision.ImageSegmenterOptions(base_options=base_options, output_category_mask=True)
-
-    with python.vision.InteractiveSegmenter.create_from_options(options) as segmenter:
-        image = mp.Image.create_from_file(image_file_path)
-        roi = vision.InteractiveSegmenterRegionOfInterest(
-            format=vision.InteractiveSegmenterRegionOfInterest.Format.KEYPOINT,
-            keypoint=containers.keypoint.NormalizedKeypoint(x, y)
-        )
-        segmentation_result = segmenter.segment(image, roi)
-        category_mask = segmentation_result.category_mask
-        segmentation_mask = category_mask.numpy_view().astype(np.uint8)
-
-        return segmentation_mask, image
-
-
 
+@spaces.GPU
 def segment(image_file_name, x, y, prompt):
     OVERLAY_COLOR = (255, 105, 180)  # Rose
 
@@ -175,6 +149,7 @@ def segment(image_file_name, x, y, prompt):
         return output_image, bbox_image
 
 
+@spaces.GPU
 def generate(image_file_path, x, y, prompt):
     output_image, bbox_image = segment(image_file_path, x, y, prompt)