wehub-resource-sync
79 lines
2.3 KiB
Python
79 lines
2.3 KiB
Python
# !pip install diffusers transformers
|
|
|
|
import requests
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
from io import BytesIO
|
|
from diffusers import DiffusionPipeline
|
|
|
|
from segment_anything import sam_model_registry, SamPredictor
|
|
|
|
|
|
"""
|
|
Step 1: Download and preprocess example demo images
|
|
"""
|
|
def download_image(url):
|
|
response = requests.get(url)
|
|
return Image.open(BytesIO(response.content)).convert("RGB")
|
|
|
|
|
|
img_url = "https://github.com/IDEA-Research/detrex-storage/blob/main/assets/grounded_sam/paint_by_example/input_image.png?raw=true"
|
|
# example_url = "https://github.com/IDEA-Research/detrex-storage/blob/main/assets/grounded_sam/paint_by_example/pomeranian_example.jpg?raw=True"
|
|
# example_url = "https://github.com/IDEA-Research/detrex-storage/blob/main/assets/grounded_sam/paint_by_example/example_image.jpg?raw=true"
|
|
example_url = "https://github.com/IDEA-Research/detrex-storage/blob/main/assets/grounded_sam/paint_by_example/labrador_example.jpg?raw=true"
|
|
|
|
init_image = download_image(img_url).resize((512, 512))
|
|
example_image = download_image(example_url).resize((512, 512))
|
|
|
|
|
|
"""
|
|
Step 2: Initialize SAM and PaintByExample models
|
|
"""
|
|
|
|
DEVICE = "cuda:1"
|
|
|
|
# SAM
|
|
SAM_ENCODER_VERSION = "vit_h"
|
|
SAM_CHECKPOINT_PATH = "/comp_robot/rentianhe/code/Grounded-Segment-Anything/sam_vit_h_4b8939.pth"
|
|
sam = sam_model_registry[SAM_ENCODER_VERSION](checkpoint=SAM_CHECKPOINT_PATH).to(device=DEVICE)
|
|
sam_predictor = SamPredictor(sam)
|
|
sam_predictor.set_image(np.array(init_image))
|
|
|
|
# PaintByExample Pipeline
|
|
CACHE_DIR = "/comp_robot/rentianhe/weights/diffusers/"
|
|
pipe = DiffusionPipeline.from_pretrained(
|
|
"Fantasy-Studio/Paint-by-Example",
|
|
torch_dtype=torch.float16,
|
|
cache_dir=CACHE_DIR,
|
|
)
|
|
pipe = pipe.to(DEVICE)
|
|
|
|
|
|
"""
|
|
Step 3: Get masks with SAM by prompt (box or point) and inpaint the mask region by example image.
|
|
"""
|
|
|
|
input_point = np.array([[350, 256]])
|
|
input_label = np.array([1]) # positive label
|
|
|
|
masks, _, _ = sam_predictor.predict(
|
|
point_coords=input_point,
|
|
point_labels=input_label,
|
|
multimask_output=False
|
|
)
|
|
mask = masks[0] # [1, 512, 512] to [512, 512] np.ndarray
|
|
mask_pil = Image.fromarray(mask)
|
|
|
|
mask_pil.save("./mask.jpg")
|
|
|
|
image = pipe(
|
|
image=init_image,
|
|
mask_image=mask_pil,
|
|
example_image=example_image,
|
|
num_inference_steps=500,
|
|
guidance_scale=9.0
|
|
).images[0]
|
|
|
|
image.save("./paint_by_example_demo.jpg")
|