Spaces:
Runtime error
Runtime error
File size: 1,681 Bytes
2e51528 4eda14d 910380b 4eda14d 2e51528 4eda14d 7542834 4eda14d 7542834 4eda14d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import spaces
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
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_vit_h.pth"
model_type = "vit_h"
device = "cuda"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
@spaces.GPU(duration=10)
def generate_image(prompt, image):
sam.to(device=device)
predictor = SamPredictor(sam)
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()')
|