m

handler.py

@@ -0,0 +1,66 @@
+import os
+import subprocess
+import torch
+import requests
+from PIL import Image
+from io import BytesIO
+
+just_get_sd_mask_function = None
+
+print(os.listdir('/usr/local/'))
+print(torch.version.cuda)
+
+class EndpointHandler():
+    def __init__(self, path="."):
+        global just_get_sd_mask_function
+        
+        is_production = True
+        
+        if False:
+            return
+
+        os.chdir(path)
+        
+        os.environ['AM_I_DOCKER'] = 'False'
+        os.environ['BUILD_WITH_CUDA'] = 'True'
+        os.environ['CUDA_HOME'] = '/usr/local/cuda-11.7/' if is_production else '/usr/local/cuda-12.1/'
+
+        # Install Segment Anything
+        subprocess.run(["python", "-m", "pip", "install", "-e", "segment_anything"])
+
+        # Install Grounding DINO
+        subprocess.run(["python", "-m", "pip", "install", "-e", "GroundingDINO"])
+    
+        subprocess.run("wget https://huggingface.co/Uminosachi/sam-hq/resolve/main/sam_hq_vit_h.pth -O ./sam_hq_vit_h.pth", shell=True)
+
+        # Install diffusers
+        subprocess.run(["pip", "install", "--upgrade", "diffusers[torch]"])
+
+        # Install osx
+        subprocess.run(["git", "submodule", "update", "--init", "--recursive"])
+        subprocess.run(["bash", "grounded-sam-osx/install.sh"], cwd="grounded-sam-osx")
+
+        # Install RAM & Tag2Text
+        subprocess.run(["git", "clone", "https://github.com/xinyu1205/recognize-anything.git"])
+        subprocess.run(["pip", "install", "-r", "./recognize-anything/requirements.txt"])
+        subprocess.run(["pip", "install", "-e", "./recognize-anything/"])
+
+        from test import just_get_sd_mask
+        just_get_sd_mask_function = just_get_sd_mask
+
+    def __call__(self, data):
+        mask_pil = just_get_sd_mask(Image.open("assets/demo1.jpg"), "bear", 10)
+
+        if mask_pil.mode != 'RGB':
+            mask_pil = mask_pil.convert('RGB')
+
+        # Convert PIL image to byte array
+        img_byte_arr = BytesIO()
+        mask_pil.save(img_byte_arr, format='JPEG')
+        img_byte_arr = img_byte_arr.getvalue()
+
+        # Upload to file.io
+        response = requests.post("https://file.io/", files={"file": img_byte_arr})
+        url = response.json().get('link')
+
+        return {"url": url}

handler_test.py

@@ -0,0 +1,13 @@
+from handler import EndpointHandler
+
+# init handler
+my_handler = EndpointHandler(path=".")
+
+# prepare sample payload
+non_holiday_payload = {"inputs": "I am quite excited how this will turn out", "date": "2022-08-08"}
+
+# test the handler
+non_holiday_pred=my_handler(non_holiday_payload)
+
+# show results
+print("non_holiday_pred", non_holiday_pred)

test.py

@@ -0,0 +1,57 @@
+from grounded_sam_demo import grounded_sam_demo
+import numpy as np
+from PIL import Image
+from scipy.ndimage import convolve
+from scipy.ndimage import binary_dilation
+
+
+def get_sd_mask(color_mask_pil, target=(72, 4, 84), tolerance=50):
+    image_array = np.array(color_mask_pil)
+
+    # Update target based on the number of color channels in the image array
+    target = np.array(list(target) + [255] *
+                      (image_array.shape[-1] - len(target)))
+
+    mask = np.abs(image_array - target) <= tolerance
+    mask = np.all(mask, axis=-1)
+
+    new_image_array = np.ones_like(image_array) * 255  # Start with white
+    # Apply black where condition met
+    new_image_array[mask] = [0] * image_array.shape[-1]
+
+    return Image.fromarray(new_image_array)
+
+
+def expand_white_pixels(input_pil, expand_by=1):
+    img_array = np.array(input_pil)
+    is_white = np.all(img_array == 255, axis=-1)
+
+    kernel = np.ones((2*expand_by+1, 2*expand_by+1), bool)
+    expanded_white = binary_dilation(is_white, structure=kernel)
+
+    expanded_array = np.where(expanded_white[..., None], 255, img_array)
+
+    expanded_pil = Image.fromarray(expanded_array.astype('uint8'))
+    return expanded_pil
+
+
+config_file = "GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py"
+grounded_checkpoint = "groundingdino_swint_ogc.pth"
+sam_checkpoint = "sam_hq_vit_h.pth"
+
+
+def just_get_sd_mask(input_pil, text_prompt, padding):
+    print("Doing sam")
+
+    colored_mask_pil = grounded_sam_demo(
+        input_pil, config_file, grounded_checkpoint, sam_checkpoint, text_prompt)
+
+    print("doing to white")
+
+    sd_mask_pil = get_sd_mask(colored_mask_pil)
+
+    print("expanding white pixels")
+
+    sd_mask_withpadding_pil = expand_white_pixels(sd_mask_pil, padding)
+
+    return sd_mask_withpadding_pil