Upload 3 files

app.py

@@ -0,0 +1,65 @@
+import gradio as gr
+import numpy as np
+import random
+import time
+import json
+import os
+from loguru import logger
+from decouple import config
+import io
+import torch
+import numpy as np
+import torch
+import cv2
+from PIL import Image
+
+from segment_anything import sam_model_registry, SamPredictor
+
+import spaces
+
+print(f"Is CUDA available: {torch.cuda.is_available()}")
+print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+print(torch.version.cuda)
+device = torch.cuda.get_device_name(torch.cuda.current_device())
+print(device)
+
+sam_checkpoint = "sam-hq/models/sam_hq_vit_h.pth"
+model_type = "vit_h"
+device = "cuda"
+sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
+sam.to(device=device)
+predictor = SamPredictor(sam)
+
+@spaces.GPU(duration=5)
+def generate_image(prompt, image):
+    predictor.set_image(image)
+
+    prompt = json.loads(prompt)
+    input_points = np.array(prompt['input_points'])
+    input_labels = np.array(prompt['input_labels'])
+
+    mask, _, _ = predictor.predict(
+        point_coords=input_points,
+        point_labels=input_labels,
+        box=None,
+        multimask_output=False,
+        hq_token_only=True,
+    )
+
+    rgb_array = np.zeros((mask.shape[1], mask.shape[2], 3), dtype=np.uint8)
+    rgb_array[mask[0]] = 255
+    result = Image.fromarray(rgb_array)
+
+    return result
+
+
+if __name__ == "__main__":
+    demo = gr.Interface(fn=generate_image, inputs=[
+        "text",
+        gr.Image(image_mode='RGB', type="numpy")
+    ],
+                        outputs=[
+                            gr.Image(type="numpy", image_mode='RGB')
+                        ])
+    demo.launch(debug=True)
+    logger.debug('demo.launch()')

requirements.txt

@@ -0,0 +1,7 @@
+segment-anything-hq
+python-decouple==3.8
+torch
+torchaudio
+torchsde
+torchvision
+loguru==0.7.2

sam_hq_vit_h.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a7ac14a085326d9fa6199c8c698c4f0e7280afdbb974d2c4660ec60877b45e35
+size 2570940653