|
|
import gc |
|
|
import gradio as gr |
|
|
from fastapi import FastAPI, Body,Header,status |
|
|
from gradio import components,Blocks,Row |
|
|
import json |
|
|
from PIL import Image |
|
|
import torchvision.transforms as transforms |
|
|
import torch |
|
|
from pathlib import Path |
|
|
import os |
|
|
import requests |
|
|
from segment_anything import sam_model_registry, SamPredictor |
|
|
import numpy as np |
|
|
from modules.safe import unsafe_torch_load, load |
|
|
from modules.devices import device, torch_gc, cpu |
|
|
|
|
|
from modules.processing import process_images |
|
|
import modules.scripts as scripts |
|
|
|
|
|
UNIT_DEBUG=False |
|
|
def import_or_install(package,pip_name=None): |
|
|
import importlib |
|
|
import subprocess |
|
|
if pip_name is None: |
|
|
pip_name=package |
|
|
try: |
|
|
importlib.import_module(package) |
|
|
print(f"{package} is already installed") |
|
|
except ImportError: |
|
|
print(f"{package} is not installed, installing now...") |
|
|
subprocess.call(['pip', 'install', package]) |
|
|
print(f"{package} has been installed") |
|
|
|
|
|
import_or_install("segment_anything","git+https://github.com/facebookresearch/segment-anything.git") |
|
|
|
|
|
class InteractiveImageSegmentor: |
|
|
def download_file_if_not_exists(file_url, file_name): |
|
|
if not os.path.isfile(file_name): |
|
|
response = requests.get(file_url) |
|
|
if response.status_code == 200: |
|
|
with open(file_name, 'wb') as file: |
|
|
file.write(response.content) |
|
|
print("File downloaded successfully!") |
|
|
else: |
|
|
print("Failed to download the file.") |
|
|
|
|
|
def load_model(self,model_choice="sam_vit_b"): |
|
|
sam_checkpoint=f"{model_choice}.pth" |
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
self.device = device |
|
|
if model_choice=="sam_vit_b":InteractiveImageSegmentor.download_file_if_not_exists("https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth",sam_checkpoint) |
|
|
elif model_choice=="sam_vit_l":InteractiveImageSegmentor.download_file_if_not_exists("https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth",sam_checkpoint) |
|
|
elif model_choice=="sam_vit_h":InteractiveImageSegmentor.download_file_if_not_exists("https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",sam_checkpoint) |
|
|
model_type=model_choice.replace("sam_","") |
|
|
if model_type not in sam_model_registry: |
|
|
model_type="default" |
|
|
print(f"Loading model {model_type} from {sam_checkpoint}") |
|
|
torch.load = unsafe_torch_load |
|
|
self.sam = sam_model_registry[model_type](checkpoint=sam_checkpoint) |
|
|
self.sam.to(self.device) |
|
|
self.predictor = SamPredictor(self.sam) |
|
|
torch.load = load |
|
|
def clear_sam_cache(self): |
|
|
self.sam.unload_model() |
|
|
gc.collect() |
|
|
torch_gc() |
|
|
|
|
|
def mask2image_multi(self,mask:torch.Tensor): |
|
|
|
|
|
if mask.dim() == 3 and mask.shape[-1] == 3: |
|
|
mask = mask.permute(2, 0, 1) |
|
|
elif mask.dim() == 3 and mask.shape[0] == 3: |
|
|
pass |
|
|
else: |
|
|
print(mask.shape) |
|
|
raise ValueError("Mask tensor has an unexpected shape.") |
|
|
color = torch.Tensor([255/255, 155/255, 114/255, 0.6]).to(self.device) |
|
|
binary_mask = mask[0, :, :] |
|
|
h, w = binary_mask.shape |
|
|
mask_image = binary_mask.reshape(h, w, 1) * color.reshape(1, 1, -1) |
|
|
return mask_image.permute(2, 0, 1) |
|
|
def mask2image(self, mask: torch.Tensor): |
|
|
if mask.dim() == 3 and mask.shape[0] == 1: |
|
|
binary_mask = mask.squeeze(0) |
|
|
elif mask.dim() == 2: |
|
|
binary_mask = mask |
|
|
else: |
|
|
print(mask.shape) |
|
|
raise ValueError("Mask tensor has an unexpected shape.") |
|
|
h, w = binary_mask.shape |
|
|
rgb_image = binary_mask.repeat(3, 1, 1) |
|
|
alpha_channel = torch.full((1, h, w), 0.6).to(self.device) |
|
|
rgba_image = torch.cat((rgb_image, alpha_channel), dim=0) |
|
|
color = torch.Tensor([255/255, 155/255, 114/255]).to(self.device).reshape(3, 1, 1) |
|
|
return rgba_image |
|
|
def preview_segment(self,image:Image,points:list[list[float]]=None,bbox=None,labels:list[int]=None): |
|
|
pil_2_tensor = transforms.PILToTensor() |
|
|
rgba_image = image.convert("RGBA") |
|
|
image_tensor = pil_2_tensor(rgba_image).cuda() |
|
|
result_tensor=image_tensor.clone() |
|
|
mask_tensor=self.segment(points,bbox,labels) |
|
|
mask_image_tensor=self.mask2image(mask_tensor) |
|
|
mask_image=transforms.ToPILImage()(mask_image_tensor) |
|
|
result_image:Image = transforms.ToPILImage()(result_tensor) |
|
|
result_image=Image.alpha_composite(result_image,mask_image) |
|
|
return result_image |
|
|
def segment(self,points:list[list[float]]=None,bbox=None,labels:list[int]=None)->torch.Tensor: |
|
|
if len(points)==0:points=None |
|
|
if len(labels)==0:labels=None |
|
|
if len(bbox)==0:bbox=None |
|
|
if points is not None:points = torch.Tensor(np.array(points)).to(self.device).unsqueeze(0) |
|
|
if labels is not None:labels = torch.Tensor(np.array(labels)).to(self.device).unsqueeze(0) |
|
|
if bbox is not None:bbox = torch.Tensor(np.array(bbox)).to(self.device) |
|
|
print(points,labels,bbox) |
|
|
masks, scores, logits = self.predictor.predict_torch( |
|
|
point_coords=points, |
|
|
point_labels=labels, |
|
|
boxes=bbox, |
|
|
multimask_output=False, |
|
|
) |
|
|
return masks[0] |
|
|
|
|
|
def remove_selected(self,image:Image,points:list[list[float]]=None,boxes=None,labels:list[int]=None): |
|
|
pil_2_tensor = transforms.PILToTensor() |
|
|
rgba_image = image.convert("RGBA") |
|
|
image_tensor = pil_2_tensor(rgba_image).cuda() |
|
|
mask_tensor = image_segmentor.segment(points=points,bbox=boxes,labels=labels) |
|
|
result_tensor=image_tensor*(1-mask_tensor) |
|
|
result_image:Image = transforms.ToPILImage()(result_tensor) |
|
|
return result_image |
|
|
def remove_unselected(self,image:Image,points:list[list[float]]=None,boxes=None,labels:list[int]=None): |
|
|
pil_2_tensor = transforms.PILToTensor() |
|
|
rgba_image = image.convert("RGBA") |
|
|
image_tensor = pil_2_tensor(rgba_image).cuda() |
|
|
mask_tensor = image_segmentor.segment(points=points,bbox=boxes,labels=labels) |
|
|
print(image_tensor.shape,mask_tensor.shape) |
|
|
result_tensor=image_tensor*mask_tensor |
|
|
result_image:Image = transforms.ToPILImage()(result_tensor) |
|
|
return result_image |
|
|
pass |
|
|
|
|
|
def reset_image(image:Image): |
|
|
global image_segmentor |
|
|
if image_segmentor is None: |
|
|
image_segmentor=InteractiveImageSegmentor() |
|
|
image_segmentor.load_model() |
|
|
image_segmentor.predictor.reset_image() |
|
|
image_array = np.array(image) |
|
|
image_segmentor.predictor.set_image(image_array) |
|
|
return image |
|
|
|
|
|
def on_image_changed(image:Image): |
|
|
global points,labels,box_cache,boxes |
|
|
points=[] |
|
|
labels=[] |
|
|
boxes=[] |
|
|
box_cache=[] |
|
|
reset_image(image) |
|
|
return image |
|
|
|
|
|
def on_image_clicked(image:Image,choice,input_type,event_data:gr.SelectData): |
|
|
global box_cache,boxes,points,labels |
|
|
if isinstance(choice,str): |
|
|
if choice=="Select":choice=1 |
|
|
elif choice=="Deselect":choice=0 |
|
|
if input_type=="Point": |
|
|
points.append(event_data.index) |
|
|
labels.append(choice) |
|
|
return image_segmentor.preview_segment(image,points=points,bbox=boxes,labels=labels) |
|
|
elif input_type=="Box": |
|
|
box_cache.extend(event_data.index) |
|
|
if len(box_cache)==4: |
|
|
boxes.append(box_cache) |
|
|
box_cache=[] |
|
|
return image_segmentor.preview_segment(image,points=points,bbox=boxes,labels=labels) |
|
|
return image |
|
|
|
|
|
def on_remove_btn_clicked(image:Image,remove_type:str): |
|
|
global points,labels,box_cache,boxes |
|
|
if remove_type=="Selected": |
|
|
return image_segmentor.remove_selected(image,points=points,boxes=boxes,labels=labels) |
|
|
elif remove_type=="Unselected": |
|
|
return image_segmentor.remove_unselected(image,points=points,boxes=boxes,labels=labels) |
|
|
return image |
|
|
|
|
|
class Script(scripts.Script): |
|
|
def title(self): |
|
|
return "Interactive Image Segmentor" |
|
|
def show(self, is_img2img): |
|
|
return is_img2img |
|
|
def ui(self, is_img2img): |
|
|
if not is_img2img: return |
|
|
with Blocks(): |
|
|
with Row(equal_height=True): |
|
|
choice=components.Radio(choices=["Select","Deselect"],value="Select",label="Selection Type") |
|
|
input_type=components.Radio(choices=["Point","Box"],value="Point",label="Input Type") |
|
|
remove_type=components.Radio(choices=["Selected","Unselected"],value="Selected",label="Remove Type") |
|
|
with Row(equal_height=True): |
|
|
image=components.Image(type="pil",interactive=True,image_mode="RGB") |
|
|
resulting_image=components.Image(type="pil",image_mode="RGBA") |
|
|
image.change(on_image_changed,inputs=[image],outputs=[resulting_image]) |
|
|
image.select(on_image_clicked,inputs=[image,choice,input_type],outputs=[resulting_image]) |
|
|
with Row(equal_height=True): |
|
|
remove_btn = components.Button(value="Preview Remove Effect") |
|
|
remove_btn.click(on_remove_btn_clicked,inputs=[image,remove_type],outputs=[resulting_image]) |
|
|
pass |
|
|
return [image,points,labels,boxes] |
|
|
|
|
|
def run(self,p,image,points,labels,boxes): |
|
|
if image is None: |
|
|
image=p.init_images[0] |
|
|
image_segmentor.predictor.set_image(np.array(image)) |
|
|
mask=image_segmentor.predictor.predict_torch(points,labels,boxes) |
|
|
p.image_mask=mask |
|
|
proc = process_images(p) |
|
|
proc.images.append(mask) |
|
|
return proc |
|
|
pass |
|
|
|
|
|
def interactive_image_segmentor_api(_: Blocks, app: FastAPI): |
|
|
@app.post("/figma/interactive_image_segmentor/upload_image") |
|
|
async def upload_image(image_str:str = Body(...)): |
|
|
import base64 |
|
|
import io |
|
|
image_bytes = base64.b64decode(image_str) |
|
|
image = Image.open(io.BytesIO(image_bytes),formats=["PNG"]) |
|
|
image_segmentor.predictor.reset_image() |
|
|
image_segmentor.predictor.set_image(np.array(image)) |
|
|
return image |
|
|
@app.post("/figma/interactive_image_segmentor/image_x_mask") |
|
|
async def remove_selected(image_str:str = Body(...),points:list[list[float]]=Body(...),\ |
|
|
boxes:list[list[float]]=Body(...),labels:list[int]=Body(...), remove_type:bool=Body(...)): |
|
|
import base64 |
|
|
import io |
|
|
image_bytes = base64.b64decode(image_str) |
|
|
image = Image.open(io.BytesIO(image_bytes),formats=["PNG"]) |
|
|
if remove_type=="Selected": |
|
|
image= image_segmentor.remove_selected(image,points=points,boxes=boxes,labels=labels) |
|
|
elif remove_type=="Unselected": |
|
|
image= image_segmentor.remove_unselected(image,points=points,boxes=boxes,labels=labels) |
|
|
return image |
|
|
pass |
|
|
|
|
|
points=[] |
|
|
labels=[] |
|
|
box_cache:list=[] |
|
|
boxes=[] |
|
|
|
|
|
image_segmentor=None |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
try: |
|
|
import modules.script_callbacks as script_callbacks |
|
|
|
|
|
script_callbacks.on_app_started(interactive_image_segmentor_api) |
|
|
except: |
|
|
pass |
|
|
pass |
