File size: 1,749 Bytes
da05695 |
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 |
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()')
|