init
This commit is contained in:
smthem
@@ -0,0 +1,42 @@
|
||||
# S-Lab License 1.0
|
||||
|
||||
Copyright 2025 S-Lab
|
||||
|
||||
Redistribution and use for non-commercial purpose in source and
|
||||
binary forms, with or without modification, are permitted provided
|
||||
that the following conditions are met:
|
||||
|
||||
1. Redistributions of source code must retain the above copyright
|
||||
notice, this list of conditions and the following disclaimer.
|
||||
|
||||
2. Redistributions in binary form must reproduce the above copyright
|
||||
notice, this list of conditions and the following disclaimer in
|
||||
the documentation and/or other materials provided with the
|
||||
distribution.
|
||||
|
||||
3. Neither the name of the copyright holder nor the names of its
|
||||
contributors may be used to endorse or promote products derived
|
||||
from this software without specific prior written permission.
|
||||
|
||||
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
|
||||
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
|
||||
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
|
||||
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
|
||||
HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
|
||||
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
|
||||
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
|
||||
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
|
||||
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
|
||||
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
|
||||
In the event that redistribution and/or use for commercial purpose in
|
||||
source or binary forms, with or without modification is required,
|
||||
please contact the contributor(s) of the work.
|
||||
|
||||
|
||||
---
|
||||
For inquiries permission for commercial use, please consult our team:
|
||||
Jixin Zhao (jixinzhao0101@gmail.com),
|
||||
Dr. Shangchen Zhou (shangchenzhou@gmail.com),
|
||||
Prof. Chen Change Loy (ccloy@ntu.edu.sg)
|
||||
@@ -0,0 +1,228 @@
|
||||
# !/usr/bin/env python
|
||||
# -*- coding: UTF-8 -*-
|
||||
import os
|
||||
import torch
|
||||
import gc
|
||||
import numpy as np
|
||||
from omegaconf import OmegaConf
|
||||
from pathlib import PureWindowsPath
|
||||
import yaml
|
||||
|
||||
from .node_utils import gc_cleanup,tensor2pil_list,load_images,mask2pil_list,tensor_upscale
|
||||
from .inference_objectclear import loader_objectclear,inference_objectclear
|
||||
import folder_paths
|
||||
|
||||
|
||||
########
|
||||
MAX_SEED = np.iinfo(np.int32).max
|
||||
current_node_path = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
device = torch.device(
|
||||
"cuda:0") if torch.cuda.is_available() else torch.device(
|
||||
"mps") if torch.backends.mps.is_available() else torch.device(
|
||||
"cpu")
|
||||
|
||||
|
||||
######
|
||||
|
||||
|
||||
class ObjectClearLoader:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"checkpoint":(["none"] + folder_paths.get_filename_list("checkpoints"),),
|
||||
"vae":(["none"] + folder_paths.get_filename_list("vae"),),
|
||||
"clip":(["none"] + folder_paths.get_filename_list("clip"),),
|
||||
"use_fp16":("BOOLEAN",{"default":True}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL_ObjectClear",)
|
||||
RETURN_NAMES = ("model",)
|
||||
FUNCTION = "loader_main"
|
||||
CATEGORY = "ObjectClear"
|
||||
|
||||
def loader_main(self,checkpoint,vae,clip,use_fp16,):
|
||||
|
||||
|
||||
# load model
|
||||
print("***********Load model ***********")
|
||||
|
||||
if clip == "none" :
|
||||
raise Exception("Please select a clip model")
|
||||
|
||||
else:
|
||||
weight_path=folder_paths.get_full_path("clip", clip)
|
||||
|
||||
if vae == "none" :
|
||||
raise Exception("Please select a vae model")
|
||||
|
||||
else:
|
||||
vae_path=folder_paths.get_full_path("vae", vae)
|
||||
|
||||
|
||||
if checkpoint == "none" :
|
||||
raise Exception("Please select a checkpoint model")
|
||||
|
||||
else:
|
||||
model_path=folder_paths.get_full_path("checkpoints", checkpoint)
|
||||
pipe = loader_objectclear(model_path,vae_path,current_node_path,device,None,use_fp16,weight_path)
|
||||
|
||||
print("***********Load model done ***********")
|
||||
|
||||
gc_cleanup()
|
||||
|
||||
return (pipe,)
|
||||
|
||||
|
||||
|
||||
class ObjectClearSampler:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("MODEL_ObjectClear",),
|
||||
"iamge": ("IMAGE",),
|
||||
"mask": ("MASK",),
|
||||
"positive": ("CONDITIONING", {
|
||||
"tooltip": "The conditioning describing the attributes you want to include in the image."}),
|
||||
"negative": ("CONDITIONING", {
|
||||
"tooltip": "The conditioning describing the attributes you want to exclude from the image."}),
|
||||
"vison_emb":("CLIP_VISION_OUTPUT",),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": MAX_SEED,
|
||||
"tooltip": "The random seed used for creating the noise."}),
|
||||
"steps": ("INT", {"default": 20, "min": 1, "max": 10000,
|
||||
"tooltip": "The number of steps used in the denoising process."}),
|
||||
"cfg": ("FLOAT", {"default": 2.5, "min": 0.1, "max": 20.0, "step": 0.1, "round": 0.01,
|
||||
"tooltip": "The Classifier-Free Guidance scale balances creativity and adherence to the prompt. Higher values result in images more closely matching the prompt however too high values will negatively impact quality."}),
|
||||
"strength": (
|
||||
"FLOAT", {"default": 0.99, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.001}),
|
||||
"short_size": (
|
||||
"INT", {"default": 512, "min": 512, "max": 2048, "step": 64,}),
|
||||
|
||||
}}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", )
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "sampler_main"
|
||||
CATEGORY = "ObjectClear"
|
||||
|
||||
def sampler_main(self,model,iamge,mask,positive,negative,vison_emb,seed,steps,cfg,strength,short_size):
|
||||
|
||||
#object_embeds=vison_emb["image_embeds"].view(1, 1, -1) # [1, 1, 768]
|
||||
B = vison_emb["image_embeds"].shape[0]
|
||||
obj_embeds = vison_emb["image_embeds"].view(B, 1, -1)
|
||||
if B==1:
|
||||
object_embeds=[obj_embeds]
|
||||
else:
|
||||
object_embeds=list(torch.chunk(obj_embeds, chunks=B))
|
||||
images_list=tensor2pil_list(iamge)
|
||||
masks_list=mask2pil_list(mask)
|
||||
images=inference_objectclear(model,images_list,masks_list,device,positive,negative,steps,seed,strength,cfg,object_embeds,short_size)
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
return (load_images(images), )
|
||||
|
||||
|
||||
class ObjectClearVision:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"iamge": ("IMAGE",),
|
||||
"mask": ("MASK",),
|
||||
}}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", )
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "sampler_main"
|
||||
CATEGORY = "ObjectClear"
|
||||
|
||||
def sampler_main(self,iamge,mask):
|
||||
B,_,_,_=iamge.size()
|
||||
B_,_,_=mask.size()
|
||||
if B!=B_:
|
||||
raise ValueError("input image and mask must have same batch size")
|
||||
else:
|
||||
if B==1:
|
||||
mask=mask.unsqueeze(-1)
|
||||
img=iamge * (mask > 0.5)
|
||||
return (img, )
|
||||
else:
|
||||
mask_list=list(torch.chunk(mask, chunks=B))
|
||||
img_list=list(torch.chunk(iamge, chunks=B))
|
||||
masked_img_list=[]
|
||||
for mask_,img_ in zip(mask_list,img_list):
|
||||
msk=mask_.unsqueeze(-1)
|
||||
image_=img_ * (msk > 0.5)
|
||||
masked_img_list.append(image_)
|
||||
return (torch.cat(masked_img_list, dim=0), )
|
||||
|
||||
class ObjectClearBatch:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"mask1": ("MASK",), # B,h,w
|
||||
"mask2": ("MASK",),
|
||||
}}
|
||||
|
||||
RETURN_TYPES = ("MASK", )
|
||||
RETURN_NAMES = ("mask",)
|
||||
FUNCTION = "sampler_main"
|
||||
CATEGORY = "ObjectClear"
|
||||
|
||||
def sampler_main(self,mask1,mask2):
|
||||
|
||||
B1,height,width=mask1.size()
|
||||
B2,_,_=mask2.size()
|
||||
|
||||
|
||||
mask1=tensor_upscale(mask1.unsqueeze(-1) ,width,height).squeeze(-1)
|
||||
mask2=tensor_upscale(mask2.unsqueeze(-1) ,width,height).squeeze(-1)
|
||||
|
||||
if B1==B2:
|
||||
return (torch.cat((mask1,mask2), dim=0), )
|
||||
else:
|
||||
out_list = []
|
||||
if B1==1:
|
||||
out_list.append(mask1)
|
||||
else:
|
||||
for i in list(torch.chunk(mask1, chunks=B1)):
|
||||
out_list.append(i)
|
||||
|
||||
if B2==1:
|
||||
out_list.append(mask2)
|
||||
else:
|
||||
for i in list(torch.chunk(mask2, chunks=B2)):
|
||||
out_list.append(i)
|
||||
return (torch.cat(out_list, dim=0), )
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
|
||||
"ObjectClearLoader": ObjectClearLoader,
|
||||
"ObjectClearSampler": ObjectClearSampler,
|
||||
"ObjectClearVision": ObjectClearVision,
|
||||
"ObjectClearBatch": ObjectClearBatch,
|
||||
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ObjectClearLoader": "ObjectClearLoader",
|
||||
"ObjectClearSampler": "ObjectClearSampler",
|
||||
"ObjectClearVision": "ObjectClearVision",
|
||||
"ObjectClearBatch": "ObjectClearBatch"
|
||||
|
||||
}
|
||||
@@ -1,16 +1,51 @@
|
||||
# ComfyUI_ObjectClear
|
||||
Complete Object Removal via Object-Effect Attention,you can try it in ComfyUI
|
||||
[ObjectClear](https://github.com/zjx0101/ObjectClear):Complete Object Removal via Object-Effect Attention,you can try it in ComfyUI
|
||||
|
||||
|
||||
# Coming soon
|
||||
# New
|
||||
* 此方法优势在于内绘不仅限于主体，比如去除物体，物体的影子也会去掉，表现优异。
|
||||
|
||||
# 1 . Installation /安装
|
||||
|
||||
In the ./ComfyUI /custom_node directory, run the following:
|
||||
```
|
||||
git clone https://github.com/smthemex/ ComfyUI_ObjectClear.git
|
||||
```
|
||||
---
|
||||
|
||||
# Example
|
||||
# 2 . Requirements
|
||||
* 通常不需要安装，因为没什么特别的库
|
||||
|
||||
```
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
# 3 . Model/模型
|
||||
* [jixin0101/ObjectClear](https://huggingface.co/jixin0101/ObjectClear/tree/main) download files as below /下载下方列出的模型:
|
||||
```
|
||||
├── your comfyUI path/models/checkpoints/
|
||||
| ├──diffusion_pytorch_model.safetensors # unet dir/ 目录，rename or not 重命名或者随你，可以下fp16的
|
||||
├── your comfyUI path/models/vae/
|
||||
| ├──diffusion_pytorch_model.safetensors # vae dir/目录，rename or not 重命名或者随你
|
||||
├── models/clip_vision/
|
||||
| ├── diffusion_pytorch_model.safetensors # image_prompt_encoder 目录，其实就是这个： clip-vit-large-patch14.safetensors
|
||||
├── models/clip/
|
||||
| ├── model.safetensors # postfuse_module dir /目录 rename or not 重命名或者随你 ，可以下fp16
|
||||
| ├── clip_l.safetensors # normal comfyUI clip 常规comfy的clip
|
||||
| ├── clip_g.safetensors # normal comfyUI clip 常规comfy的clip
|
||||
|
||||
```
|
||||
* seg /遮罩没有内置，随便用其他的吧。
|
||||
|
||||
# 4 . Example
|
||||
|
||||
|
||||
|
||||
|
||||
# 5 .Tips/使用说明
|
||||
* 支持单图或多图，遮罩合并用插件带的batch节点，512是垫图最小短边尺寸（会自动裁切），模型用512*512训练的，按理此尺寸是最好的。
|
||||
* Support single or multiple images, use batch nodes with plugins for mask merging, 512 is the input image ‘s minimum short edge size, and the model is trained with 512 * 512, which is theoretically the best size.
|
||||
|
||||
# Citation
|
||||
```
|
||||
@InProceedings{zhao2025ObjectClear,
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
|
||||
from .ObjectClear_node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
|
Before Width: | Height: | Size: 202 KiB After Width: | Height: | Size: 202 KiB |
|
Before Width: | Height: | Size: 462 KiB After Width: | Height: | Size: 462 KiB |
@@ -0,0 +1,832 @@
|
||||
{
|
||||
"id": "d49a2aaa-7f48-41e7-b5e8-24619dec06fd",
|
||||
"revision": 0,
|
||||
"last_node_id": 30,
|
||||
"last_link_id": 46,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 9,
|
||||
"type": "CLIPTextEncode",
|
||||
"pos": [
|
||||
20464.234375,
|
||||
-950.14892578125
|
||||
],
|
||||
"size": [
|
||||
400,
|
||||
200
|
||||
],
|
||||
"flags": {
|
||||
"collapsed": true
|
||||
},
|
||||
"order": 8,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "clip",
|
||||
"type": "CLIP",
|
||||
"link": 19
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "CONDITIONING",
|
||||
"type": "CONDITIONING",
|
||||
"links": [
|
||||
25
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "CLIPTextEncode"
|
||||
},
|
||||
"widgets_values": [
|
||||
""
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 8,
|
||||
"type": "CLIPTextEncode",
|
||||
"pos": [
|
||||
20460.73046875,
|
||||
-991.8788452148438
|
||||
],
|
||||
"size": [
|
||||
400,
|
||||
200
|
||||
],
|
||||
"flags": {
|
||||
"collapsed": true
|
||||
},
|
||||
"order": 7,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "clip",
|
||||
"type": "CLIP",
|
||||
"link": 18
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "CONDITIONING",
|
||||
"type": "CONDITIONING",
|
||||
"links": [
|
||||
24
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "CLIPTextEncode"
|
||||
},
|
||||
"widgets_values": [
|
||||
"remove the instance of object"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 14,
|
||||
"type": "ObjectClearLoader",
|
||||
"pos": [
|
||||
20623.1328125,
|
||||
-1182.1590576171875
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
130
|
||||
],
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "model",
|
||||
"type": "MODEL_ObjectClear",
|
||||
"links": [
|
||||
21
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "ObjectClearLoader"
|
||||
},
|
||||
"widgets_values": [
|
||||
"diffusion_pytorch_model.fp16.safetensors",
|
||||
"diffusion_pytorch_model.safetensors",
|
||||
"postfuse_module.safetensors",
|
||||
true
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 15,
|
||||
"type": "DualCLIPLoader",
|
||||
"pos": [
|
||||
20056.70703125,
|
||||
-1170.5067138671875
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
130
|
||||
],
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "CLIP",
|
||||
"type": "CLIP",
|
||||
"links": [
|
||||
18,
|
||||
19
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "DualCLIPLoader"
|
||||
},
|
||||
"widgets_values": [
|
||||
"clip_g.safetensors",
|
||||
"clip_l.safetensors",
|
||||
"sdxl",
|
||||
"default"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 17,
|
||||
"type": "CLIPVisionLoader",
|
||||
"pos": [
|
||||
20055.10546875,
|
||||
-984.60302734375
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
58
|
||||
],
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "CLIP_VISION",
|
||||
"type": "CLIP_VISION",
|
||||
"links": [
|
||||
27
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "CLIPVisionLoader"
|
||||
},
|
||||
"widgets_values": [
|
||||
"clip-vit-large-patch14.safetensors"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 16,
|
||||
"type": "ObjectClearSampler",
|
||||
"pos": [
|
||||
20882.8359375,
|
||||
-956.9453125
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
278
|
||||
],
|
||||
"flags": {},
|
||||
"order": 15,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "model",
|
||||
"type": "MODEL_ObjectClear",
|
||||
"link": 21
|
||||
},
|
||||
{
|
||||
"name": "iamge",
|
||||
"type": "IMAGE",
|
||||
"link": 41
|
||||
},
|
||||
{
|
||||
"name": "mask",
|
||||
"type": "MASK",
|
||||
"link": 46
|
||||
},
|
||||
{
|
||||
"name": "positive",
|
||||
"type": "CONDITIONING",
|
||||
"link": 24
|
||||
},
|
||||
{
|
||||
"name": "negative",
|
||||
"type": "CONDITIONING",
|
||||
"link": 25
|
||||
},
|
||||
{
|
||||
"name": "vison_emb",
|
||||
"type": "CLIP_VISION_OUTPUT",
|
||||
"link": 34
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
26
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "ObjectClearSampler"
|
||||
},
|
||||
"widgets_values": [
|
||||
1063207822,
|
||||
"randomize",
|
||||
20,
|
||||
2.5,
|
||||
0.99,
|
||||
512
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 28,
|
||||
"type": "ImageBatch",
|
||||
"pos": [
|
||||
20383.22265625,
|
||||
-852.30224609375
|
||||
],
|
||||
"size": [
|
||||
140,
|
||||
46
|
||||
],
|
||||
"flags": {},
|
||||
"order": 11,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "image1",
|
||||
"type": "IMAGE",
|
||||
"link": 39
|
||||
},
|
||||
{
|
||||
"name": "image2",
|
||||
"type": "IMAGE",
|
||||
"link": 40
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
41,
|
||||
42
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "ImageBatch"
|
||||
},
|
||||
"widgets_values": []
|
||||
},
|
||||
{
|
||||
"id": 25,
|
||||
"type": "ImageToMask",
|
||||
"pos": [
|
||||
20122.365234375,
|
||||
-383.7356262207031
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
58
|
||||
],
|
||||
"flags": {},
|
||||
"order": 10,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"link": 38
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": [
|
||||
43
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "ImageToMask"
|
||||
},
|
||||
"widgets_values": [
|
||||
"red"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 5,
|
||||
"type": "ImageToMask",
|
||||
"pos": [
|
||||
20095.9609375,
|
||||
-191.02154541015625
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
58
|
||||
],
|
||||
"flags": {},
|
||||
"order": 9,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"link": 3
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": [
|
||||
44
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "ImageToMask"
|
||||
},
|
||||
"widgets_values": [
|
||||
"red"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 30,
|
||||
"type": "ObjectClearBatch",
|
||||
"pos": [
|
||||
20397.666015625,
|
||||
-665.2557373046875
|
||||
],
|
||||
"size": [
|
||||
149.61679077148438,
|
||||
46
|
||||
],
|
||||
"flags": {},
|
||||
"order": 12,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mask1",
|
||||
"type": "MASK",
|
||||
"link": 43
|
||||
},
|
||||
{
|
||||
"name": "mask2",
|
||||
"type": "MASK",
|
||||
"link": 44
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "mask",
|
||||
"type": "MASK",
|
||||
"links": [
|
||||
45,
|
||||
46
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "ObjectClearBatch"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 24,
|
||||
"type": "ObjectClearVision",
|
||||
"pos": [
|
||||
20683.552734375,
|
||||
-532.195068359375
|
||||
],
|
||||
"size": [
|
||||
151.69491577148438,
|
||||
46
|
||||
],
|
||||
"flags": {},
|
||||
"order": 13,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "iamge",
|
||||
"type": "IMAGE",
|
||||
"link": 42
|
||||
},
|
||||
{
|
||||
"name": "mask",
|
||||
"type": "MASK",
|
||||
"link": 45
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
35
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "ObjectClearVision"
|
||||
},
|
||||
"widgets_values": []
|
||||
},
|
||||
{
|
||||
"id": 18,
|
||||
"type": "CLIPVisionEncode",
|
||||
"pos": [
|
||||
20937.95703125,
|
||||
-610.2886352539062
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
78
|
||||
],
|
||||
"flags": {},
|
||||
"order": 14,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "clip_vision",
|
||||
"type": "CLIP_VISION",
|
||||
"link": 27
|
||||
},
|
||||
{
|
||||
"name": "image",
|
||||
"type": "IMAGE",
|
||||
"link": 35
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "CLIP_VISION_OUTPUT",
|
||||
"type": "CLIP_VISION_OUTPUT",
|
||||
"links": [
|
||||
34
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "CLIPVisionEncode"
|
||||
},
|
||||
"widgets_values": [
|
||||
"none"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 6,
|
||||
"type": "SaveImage",
|
||||
"pos": [
|
||||
20545.396484375,
|
||||
-420.887451171875
|
||||
],
|
||||
"size": [
|
||||
796.6367797851562,
|
||||
440.71405029296875
|
||||
],
|
||||
"flags": {},
|
||||
"order": 16,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 26
|
||||
}
|
||||
],
|
||||
"outputs": [],
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
"ComfyUI"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 4,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
19725.62109375,
|
||||
-88.6943130493164
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
314
|
||||
],
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
3
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"test-sample5.png",
|
||||
"image"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 26,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
19704.015625,
|
||||
-475.34197998046875
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
314
|
||||
],
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
38
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"test-sample4.png",
|
||||
"image"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 3,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
19944.125,
|
||||
-842.93603515625
|
||||
],
|
||||
"size": [
|
||||
369.85162353515625,
|
||||
326.88409423828125
|
||||
],
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
40
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"test-sample5.jpg",
|
||||
"image"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 27,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
19555.708984375,
|
||||
-877.9942016601562
|
||||
],
|
||||
"size": [
|
||||
350.5255126953125,
|
||||
326
|
||||
],
|
||||
"flags": {},
|
||||
"order": 6,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
39
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"test-sample4.jpg",
|
||||
"image"
|
||||
]
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
3,
|
||||
4,
|
||||
0,
|
||||
5,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
18,
|
||||
15,
|
||||
0,
|
||||
8,
|
||||
0,
|
||||
"CLIP"
|
||||
],
|
||||
[
|
||||
19,
|
||||
15,
|
||||
0,
|
||||
9,
|
||||
0,
|
||||
"CLIP"
|
||||
],
|
||||
[
|
||||
21,
|
||||
14,
|
||||
0,
|
||||
16,
|
||||
0,
|
||||
"MODEL_ObjectClear"
|
||||
],
|
||||
[
|
||||
24,
|
||||
8,
|
||||
0,
|
||||
16,
|
||||
3,
|
||||
"CONDITIONING"
|
||||
],
|
||||
[
|
||||
25,
|
||||
9,
|
||||
0,
|
||||
16,
|
||||
4,
|
||||
"CONDITIONING"
|
||||
],
|
||||
[
|
||||
26,
|
||||
16,
|
||||
0,
|
||||
6,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
27,
|
||||
17,
|
||||
0,
|
||||
18,
|
||||
0,
|
||||
"CLIP_VISION"
|
||||
],
|
||||
[
|
||||
34,
|
||||
18,
|
||||
0,
|
||||
16,
|
||||
5,
|
||||
"CLIP_VISION_OUTPUT"
|
||||
],
|
||||
[
|
||||
35,
|
||||
24,
|
||||
0,
|
||||
18,
|
||||
1,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
38,
|
||||
26,
|
||||
0,
|
||||
25,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
39,
|
||||
27,
|
||||
0,
|
||||
28,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
40,
|
||||
3,
|
||||
0,
|
||||
28,
|
||||
1,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
41,
|
||||
28,
|
||||
0,
|
||||
16,
|
||||
1,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
42,
|
||||
28,
|
||||
0,
|
||||
24,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
43,
|
||||
25,
|
||||
0,
|
||||
30,
|
||||
0,
|
||||
"MASK"
|
||||
],
|
||||
[
|
||||
44,
|
||||
5,
|
||||
0,
|
||||
30,
|
||||
1,
|
||||
"MASK"
|
||||
],
|
||||
[
|
||||
45,
|
||||
30,
|
||||
0,
|
||||
24,
|
||||
1,
|
||||
"MASK"
|
||||
],
|
||||
[
|
||||
46,
|
||||
30,
|
||||
0,
|
||||
16,
|
||||
2,
|
||||
"MASK"
|
||||
]
|
||||
],
|
||||
"groups": [],
|
||||
"config": {},
|
||||
"extra": {
|
||||
"ds": {
|
||||
"scale": 0.6209213230591554,
|
||||
"offset": [
|
||||
-19288.882829214504,
|
||||
1249.3984719373475
|
||||
]
|
||||
},
|
||||
"frontendVersion": "1.23.4"
|
||||
},
|
||||
"version": 0.4
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 51 KiB |
@@ -0,0 +1,657 @@
|
||||
import sys
|
||||
sys.path.append("../")
|
||||
|
||||
import gradio as gr
|
||||
import os
|
||||
from PIL import Image
|
||||
import torch
|
||||
from diffusers.utils import check_min_version
|
||||
from tools.download_util import load_file_from_url
|
||||
from tools.painter import mask_painter
|
||||
import argparse
|
||||
from objectclear.pipelines import ObjectClearPipeline
|
||||
import numpy as np
|
||||
import torchvision.transforms.functional as TF
|
||||
from scipy.ndimage import convolve, zoom
|
||||
import cv2
|
||||
import time
|
||||
|
||||
from tools.interact_tools import SamControler
|
||||
from tools.misc import get_device
|
||||
import json
|
||||
|
||||
check_min_version("0.30.2")
|
||||
|
||||
|
||||
def parse_augment():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--device', type=str, default=None)
|
||||
parser.add_argument('--sam_model_type', type=str, default="vit_h")
|
||||
parser.add_argument('--port', type=int, default=8000, help="only useful when running gradio applications")
|
||||
args = parser.parse_args()
|
||||
|
||||
if not args.device:
|
||||
args.device = str(get_device())
|
||||
|
||||
return args
|
||||
|
||||
|
||||
def pad_to_multiple(image: np.ndarray, multiple: int = 8):
|
||||
h, w = image.shape[:2]
|
||||
pad_h = (multiple - h % multiple) % multiple
|
||||
pad_w = (multiple - w % multiple) % multiple
|
||||
if image.ndim == 3:
|
||||
padded = np.pad(image, ((0, pad_h), (0, pad_w), (0,0)), mode='reflect')
|
||||
else:
|
||||
padded = np.pad(image, ((0, pad_h), (0, pad_w)), mode='reflect')
|
||||
return padded, h, w
|
||||
|
||||
def crop_to_original(image: np.ndarray, h: int, w: int):
|
||||
return image[:h, :w]
|
||||
|
||||
def wavelet_blur_np(image: np.ndarray, radius: int):
|
||||
kernel = np.array([
|
||||
[0.0625, 0.125, 0.0625],
|
||||
[0.125, 0.25, 0.125],
|
||||
[0.0625, 0.125, 0.0625]
|
||||
], dtype=np.float32)
|
||||
|
||||
blurred = np.empty_like(image)
|
||||
for c in range(image.shape[0]):
|
||||
blurred_c = convolve(image[c], kernel, mode='nearest')
|
||||
if radius > 1:
|
||||
blurred_c = zoom(zoom(blurred_c, 1 / radius, order=1), radius, order=1)
|
||||
blurred[c] = blurred_c
|
||||
return blurred
|
||||
|
||||
def wavelet_decomposition_np(image: np.ndarray, levels=5):
|
||||
high_freq = np.zeros_like(image)
|
||||
for i in range(levels):
|
||||
radius = 2 ** i
|
||||
low_freq = wavelet_blur_np(image, radius)
|
||||
high_freq += (image - low_freq)
|
||||
image = low_freq
|
||||
return high_freq, low_freq
|
||||
|
||||
def wavelet_reconstruction_np(content_feat: np.ndarray, style_feat: np.ndarray):
|
||||
content_high, _ = wavelet_decomposition_np(content_feat)
|
||||
_, style_low = wavelet_decomposition_np(style_feat)
|
||||
return content_high + style_low
|
||||
|
||||
def wavelet_color_fix_np(fused: np.ndarray, mask: np.ndarray) -> np.ndarray:
|
||||
fused_np = fused.astype(np.float32) / 255.0
|
||||
mask_np = mask.astype(np.float32) / 255.0
|
||||
|
||||
fused_np = fused_np.transpose(2, 0, 1)
|
||||
mask_np = mask_np.transpose(2, 0, 1)
|
||||
|
||||
result_np = wavelet_reconstruction_np(fused_np, mask_np)
|
||||
|
||||
result_np = result_np.transpose(1, 2, 0)
|
||||
result_np = np.clip(result_np * 255.0, 0, 255).astype(np.uint8)
|
||||
|
||||
return result_np
|
||||
|
||||
def fuse_with_wavelet(ori: np.ndarray, removed: np.ndarray, attn_map: np.ndarray, multiple: int = 8):
|
||||
H, W = ori.shape[:2]
|
||||
attn_map = attn_map.astype(np.float32)
|
||||
_, attn_map = cv2.threshold(attn_map, 128, 255, cv2.THRESH_BINARY)
|
||||
am = attn_map.astype(np.float32)
|
||||
am = am/255.0
|
||||
am_up = cv2.resize(am, (W, H), interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (21,21))
|
||||
am_d = cv2.dilate(am_up, kernel, iterations=1)
|
||||
am_d = cv2.GaussianBlur(am_d.astype(np.float32), (9,9), sigmaX=2)
|
||||
|
||||
am_merged = np.maximum(am_up, am_d)
|
||||
am_merged = np.clip(am_merged, 0, 1)
|
||||
|
||||
attn_up_3c = np.stack([am_merged]*3, axis=-1)
|
||||
attn_up_ori_3c = np.stack([am_up]*3, axis=-1)
|
||||
|
||||
ori_out = ori * (1 - attn_up_ori_3c)
|
||||
rem_out = removed * (1 - attn_up_ori_3c)
|
||||
|
||||
ori_pad, h0, w0 = pad_to_multiple(ori_out, multiple)
|
||||
rem_pad, _, _ = pad_to_multiple(rem_out, multiple)
|
||||
|
||||
wave_rgb = wavelet_color_fix_np(ori_pad, rem_pad)
|
||||
wave = crop_to_original(wave_rgb, h0, w0)
|
||||
# fusion
|
||||
fused = (wave * (1 - attn_up_3c) + removed * attn_up_3c).astype(np.uint8)
|
||||
return fused
|
||||
|
||||
|
||||
def resize_by_short_side(image, target_short=512, resample=Image.BICUBIC):
|
||||
w, h = image.size
|
||||
if w < h:
|
||||
new_w = target_short
|
||||
new_h = int(h * target_short / w)
|
||||
new_h = (new_h + 15) // 16 * 16
|
||||
else:
|
||||
new_h = target_short
|
||||
new_w = int(w * target_short / h)
|
||||
new_w = (new_w + 15) // 16 * 16
|
||||
return image.resize((new_w, new_h), resample=resample)
|
||||
|
||||
# convert points input to prompt state
|
||||
def get_prompt(click_state, click_input):
|
||||
inputs = json.loads(click_input)
|
||||
points = click_state[0]
|
||||
labels = click_state[1]
|
||||
for input in inputs:
|
||||
points.append(input[:2])
|
||||
labels.append(input[2])
|
||||
click_state[0] = points
|
||||
click_state[1] = labels
|
||||
prompt = {
|
||||
"prompt_type":["click"],
|
||||
"input_point":click_state[0],
|
||||
"input_label":click_state[1],
|
||||
"multimask_output":"True",
|
||||
}
|
||||
return prompt
|
||||
|
||||
# use sam to get the mask
|
||||
def sam_refine(image_state, point_prompt, click_state, evt:gr.SelectData):
|
||||
if point_prompt == "Positive":
|
||||
coordinate = "[[{},{},1]]".format(evt.index[0], evt.index[1])
|
||||
else:
|
||||
coordinate = "[[{},{},0]]".format(evt.index[0], evt.index[1])
|
||||
|
||||
# prompt for sam model
|
||||
model.samcontroler.sam_controler.reset_image()
|
||||
model.samcontroler.sam_controler.set_image(image_state["origin_image"])
|
||||
prompt = get_prompt(click_state=click_state, click_input=coordinate)
|
||||
|
||||
mask, logit, painted_image = model.first_frame_click(
|
||||
image=image_state["origin_image"],
|
||||
points=np.array(prompt["input_point"]),
|
||||
labels=np.array(prompt["input_label"]),
|
||||
multimask=prompt["multimask_output"],
|
||||
)
|
||||
image_state["mask"] = mask
|
||||
image_state["logit"] = logit
|
||||
image_state["painted_image"] = painted_image
|
||||
|
||||
return painted_image, image_state, click_state
|
||||
|
||||
|
||||
def add_multi_mask(image_state, interactive_state, mask_dropdown):
|
||||
mask = image_state["mask"]
|
||||
interactive_state["masks"].append(mask)
|
||||
interactive_state["mask_names"].append("mask_{:03d}".format(len(interactive_state["masks"])))
|
||||
mask_dropdown.append("mask_{:03d}".format(len(interactive_state["masks"])))
|
||||
select_frame = show_mask(image_state, interactive_state, mask_dropdown)
|
||||
|
||||
return interactive_state, gr.update(choices=interactive_state["mask_names"], value=mask_dropdown), select_frame, [[],[]]
|
||||
|
||||
def clear_click(image_state, click_state):
|
||||
click_state = [[],[]]
|
||||
input_image = image_state["origin_image"]
|
||||
return input_image, click_state
|
||||
|
||||
def remove_multi_mask(interactive_state, click_state, image_state):
|
||||
interactive_state["mask_names"]= []
|
||||
interactive_state["masks"] = []
|
||||
click_state = [[],[]]
|
||||
input_image = image_state["origin_image"]
|
||||
|
||||
return interactive_state, gr.update(choices=[],value=[]), input_image, click_state
|
||||
|
||||
def show_mask(image_state, interactive_state, mask_dropdown):
|
||||
mask_dropdown.sort()
|
||||
if image_state["origin_image"] is not None:
|
||||
select_frame = image_state["origin_image"]
|
||||
for i in range(len(mask_dropdown)):
|
||||
mask_number = int(mask_dropdown[i].split("_")[1]) - 1
|
||||
mask = interactive_state["masks"][mask_number]
|
||||
select_frame = mask_painter(select_frame, mask.astype('uint8'), mask_color=mask_number+2)
|
||||
|
||||
return select_frame
|
||||
|
||||
def upload_and_reset(image_input, interactive_state):
|
||||
click_state = [[], []]
|
||||
|
||||
interactive_state["mask_names"]= []
|
||||
interactive_state["masks"] = []
|
||||
|
||||
image_state, image_info, image_input = update_image_state_on_upload(image_input)
|
||||
|
||||
return (
|
||||
image_state,
|
||||
image_info,
|
||||
image_input,
|
||||
interactive_state,
|
||||
click_state,
|
||||
gr.update(choices=[], value=[]),
|
||||
)
|
||||
|
||||
def update_image_state_on_upload(image_input):
|
||||
frame = image_input
|
||||
|
||||
image_size = (frame.size[1], frame.size[0])
|
||||
|
||||
frame_np = np.array(frame)
|
||||
|
||||
image_state = {
|
||||
"origin_image": frame_np,
|
||||
"painted_image": frame_np.copy(),
|
||||
"mask": np.zeros((image_size[0], image_size[1]), np.uint8),
|
||||
"logit": None,
|
||||
}
|
||||
|
||||
image_info = f"Image Name: uploaded.png,\nImage Size: {image_size}"
|
||||
|
||||
model.samcontroler.sam_controler.reset_image()
|
||||
model.samcontroler.sam_controler.set_image(frame_np)
|
||||
|
||||
return image_state, image_info, image_input
|
||||
|
||||
|
||||
# SAM generator
|
||||
class MaskGenerator():
|
||||
def __init__(self, sam_checkpoint, args):
|
||||
self.args = args
|
||||
self.samcontroler = SamControler(sam_checkpoint, args.sam_model_type, args.device)
|
||||
|
||||
def first_frame_click(self, image: np.ndarray, points:np.ndarray, labels: np.ndarray, multimask=True):
|
||||
mask, logit, painted_image = self.samcontroler.first_frame_click(image, points, labels, multimask)
|
||||
return mask, logit, painted_image
|
||||
|
||||
|
||||
# args, defined in track_anything.py
|
||||
args = parse_augment()
|
||||
sam_checkpoint_url_dict = {
|
||||
'vit_h': "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",
|
||||
'vit_l': "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth",
|
||||
'vit_b': "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
|
||||
}
|
||||
|
||||
sam_checkpoint = load_file_from_url(sam_checkpoint_url_dict[args.sam_model_type])
|
||||
# initialize sams
|
||||
model = MaskGenerator(sam_checkpoint, args)
|
||||
|
||||
# Build pipeline
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
pipe = ObjectClearPipeline.from_pretrained_with_custom_modules(
|
||||
"jixin0101/ObjectClear",
|
||||
torch_dtype=torch.float16,
|
||||
variant='fp16',
|
||||
apply_attention_guided_fusion=True,
|
||||
cache_dir=None,
|
||||
)
|
||||
|
||||
pipe.to(device)
|
||||
|
||||
def process(image_state, interactive_state, mask_dropdown, guidance_scale, seed, num_inference_steps, strength
|
||||
):
|
||||
generator = torch.Generator(device="cuda").manual_seed(seed)
|
||||
image_np = image_state["origin_image"]
|
||||
image = Image.fromarray(image_np)
|
||||
if interactive_state["masks"]:
|
||||
if len(mask_dropdown) == 0:
|
||||
mask_dropdown = ["mask_001"]
|
||||
mask_dropdown.sort()
|
||||
template_mask = interactive_state["masks"][int(mask_dropdown[0].split("_")[1]) - 1] * (int(mask_dropdown[0].split("_")[1]))
|
||||
for i in range(1,len(mask_dropdown)):
|
||||
mask_number = int(mask_dropdown[i].split("_")[1]) - 1
|
||||
template_mask = np.clip(template_mask+interactive_state["masks"][mask_number]*(mask_number+1), 0, mask_number+1)
|
||||
image_state["mask"]= template_mask
|
||||
else:
|
||||
template_mask = image_state["mask"]
|
||||
mask = Image.fromarray((template_mask).astype(np.uint8) * 255)
|
||||
image_or = image.copy()
|
||||
|
||||
image = image.convert("RGB")
|
||||
mask = mask.convert("RGB")
|
||||
|
||||
image = resize_by_short_side(image, 512, resample=Image.BICUBIC)
|
||||
mask = resize_by_short_side(mask, 512, resample=Image.NEAREST)
|
||||
|
||||
w, h = image.size
|
||||
|
||||
result = pipe(
|
||||
prompt="remove the instance of object",
|
||||
image=image,
|
||||
mask_image=mask,
|
||||
generator=generator,
|
||||
num_inference_steps=num_inference_steps,
|
||||
strength=strength,
|
||||
guidance_scale=guidance_scale,
|
||||
height=h,
|
||||
width=w,
|
||||
)
|
||||
fused_img_pil = result.images[0]
|
||||
|
||||
return fused_img_pil.resize((image_or.size[:2])), (image.resize((image_or.size[:2])), fused_img_pil.resize((image_or.size[:2])))
|
||||
|
||||
import base64
|
||||
with open("./Logo.png", "rb") as f:
|
||||
img_bytes = f.read()
|
||||
img_b64 = base64.b64encode(img_bytes).decode()
|
||||
|
||||
html_img = f'''
|
||||
<div style="display:flex; justify-content:center; align-items:center; width:100%;">
|
||||
<img src="data:image/png;base64,{img_b64}" style="border:none; width:200px; height:auto;"/>
|
||||
</div>
|
||||
'''
|
||||
|
||||
tutorial_url = "https://github.com/zjx0101/ObjectClear/releases/download/media/tutorial.mp4"
|
||||
assets_path = './'
|
||||
load_file_from_url(tutorial_url, assets_path)
|
||||
|
||||
description = r"""
|
||||
<b>Official Gradio demo</b> for <a href='https://github.com/zjx0101/ObjectClear' target='_blank'><b>ObjectClear: Complete Object Removal via Object-Effect Attention</b></a>.<br>
|
||||
🔥 ObjectClear is an object removal model that can jointly eliminate the target object and its associated effects leveraging Object-Effect Attention, while preserving background consistency.<br>
|
||||
🖼️ Try to drop your image, assign the target masks with a few clicks, and get the object removal results!<br>
|
||||
|
||||
*Note: All input images are temporarily resized (shorter side = 512 pixels) during inference to match the training resolution. Final outputs are restored to the original resolution.<br>*
|
||||
"""
|
||||
|
||||
article = r"""<h3>
|
||||
<b>If ObjectClear is helpful, please help to star the <a href='https://github.com/zjx0101/ObjectClear' target='_blank'>Github Repo</a>. Thanks!</b></h3>
|
||||
<hr>
|
||||
|
||||
📑 **Citation**
|
||||
<br>
|
||||
If our work is useful for your research, please consider citing:
|
||||
```bibtex
|
||||
@InProceedings{zhao2025ObjectClear,
|
||||
title = {{ObjectClear}: Complete Object Removal via Object-Effect Attention},
|
||||
author = {Zhao, Jixin and Zhou, Shangchen and Wang, Zhouxia and Yang, Peiqing and Loy, Chen Change},
|
||||
booktitle = {arXiv preprint arXiv:2505.22636},
|
||||
year = {2025}
|
||||
}
|
||||
```
|
||||
📧 **Contact**
|
||||
<br>
|
||||
If you have any questions, please feel free to reach me out at <b>jixinzhao0101@gmail.com</b>.
|
||||
<br>
|
||||
👏 **Acknowledgement**
|
||||
<br>
|
||||
This demo is adapted from [MatAnyone](https://github.com/pq-yang/MatAnyone), and leveraging segmentation capabilities from [Segment Anything](https://github.com/facebookresearch/segment-anything). Thanks for their awesome works!
|
||||
"""
|
||||
|
||||
custom_css = """
|
||||
#input-image {
|
||||
aspect-ratio: 1 / 1;
|
||||
width: 100%;
|
||||
max-width: 100%;
|
||||
height: auto;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
}
|
||||
|
||||
#input-image img {
|
||||
max-width: 100%;
|
||||
max-height: 100%;
|
||||
object-fit: contain;
|
||||
display: block;
|
||||
}
|
||||
|
||||
#main-columns {
|
||||
gap: 60px;
|
||||
}
|
||||
|
||||
#main-columns > .gr-column {
|
||||
flex: 1;
|
||||
}
|
||||
|
||||
#compare-image {
|
||||
width: 100%;
|
||||
aspect-ratio: 1 / 1;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
margin: 0;
|
||||
padding: 0;
|
||||
max-width: 100%;
|
||||
box-sizing: border-box;
|
||||
}
|
||||
|
||||
#compare-image svg.svelte-zyxd38 {
|
||||
position: absolute !important;
|
||||
top: 50% !important;
|
||||
left: 50% !important;
|
||||
transform: translate(-50%, -50%) !important;
|
||||
}
|
||||
|
||||
#compare-image .icon.svelte-1oiin9d {
|
||||
position: absolute;
|
||||
top: 50%;
|
||||
left: 50%;
|
||||
transform: translate(-50%, -50%);
|
||||
}
|
||||
|
||||
#compare-image {
|
||||
position: relative;
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
.new_button {background-color: #171717 !important; color: #ffffff !important; border: none !important;}
|
||||
.new_button:hover {background-color: #4b4b4b !important;}
|
||||
|
||||
#start-button {
|
||||
background: linear-gradient(135deg, #2575fc 0%, #6a11cb 100%);
|
||||
color: white;
|
||||
border: none;
|
||||
padding: 12px 24px;
|
||||
font-size: 16px;
|
||||
font-weight: bold;
|
||||
border-radius: 12px;
|
||||
cursor: pointer;
|
||||
box-shadow: 0 0 12px rgba(100, 100, 255, 0.7);
|
||||
transition: all 0.3s ease;
|
||||
}
|
||||
#start-button:hover {
|
||||
transform: scale(1.05);
|
||||
box-shadow: 0 0 20px rgba(100, 100, 255, 1);
|
||||
}
|
||||
|
||||
<style>
|
||||
.button-wrapper {
|
||||
width: 30%;
|
||||
text-align: center;
|
||||
}
|
||||
.wide-button {
|
||||
width: 83% !important;
|
||||
background-color: black !important;
|
||||
color: white !important;
|
||||
border: none !important;
|
||||
padding: 8px 0 !important;
|
||||
font-size: 16px !important;
|
||||
display: inline-block;
|
||||
margin: 30px 0px 0px 50px ;
|
||||
}
|
||||
.wide-button:hover {
|
||||
background-color: #656262 !important;
|
||||
}
|
||||
</style>
|
||||
"""
|
||||
|
||||
|
||||
with gr.Blocks(css=custom_css) as demo:
|
||||
gr.HTML(html_img)
|
||||
gr.Markdown(description)
|
||||
with gr.Group(elem_classes="gr-monochrome-group", visible=True):
|
||||
with gr.Row():
|
||||
with gr.Accordion('SAM Settings (click to expand)', open=False):
|
||||
with gr.Row():
|
||||
point_prompt = gr.Radio(
|
||||
choices=["Positive", "Negative"],
|
||||
value="Positive",
|
||||
label="Point Prompt",
|
||||
info="Click to add positive or negative point for target mask",
|
||||
interactive=True,
|
||||
min_width=100,
|
||||
scale=1)
|
||||
mask_dropdown = gr.Dropdown(multiselect=True, value=[], label="Mask Selection", info="Choose 1~all mask(s) added in Step 2")
|
||||
|
||||
with gr.Row(elem_id="main-columns"):
|
||||
with gr.Column():
|
||||
|
||||
click_state = gr.State([[],[]])
|
||||
|
||||
interactive_state = gr.State(
|
||||
{
|
||||
"mask_names": [],
|
||||
"masks": []
|
||||
}
|
||||
)
|
||||
|
||||
image_state = gr.State(
|
||||
{
|
||||
"origin_image": None,
|
||||
"painted_image": None,
|
||||
"mask": None,
|
||||
"logit": None
|
||||
}
|
||||
)
|
||||
|
||||
image_info = gr.Textbox(label="Image Info", visible=False)
|
||||
input_image = gr.Image(
|
||||
label='Input',
|
||||
type='pil',
|
||||
sources=["upload"],
|
||||
image_mode='RGB',
|
||||
interactive=True,
|
||||
elem_id="input-image"
|
||||
)
|
||||
|
||||
with gr.Row(equal_height=True, elem_classes="mask_button_group"):
|
||||
clear_button_click = gr.Button(value="Clear Clicks",elem_classes="new_button", min_width=100)
|
||||
add_mask_button = gr.Button(value="Add Mask", elem_classes="new_button", min_width=100)
|
||||
remove_mask_button = gr.Button(value="Delete Mask", elem_classes="new_button", min_width=100)
|
||||
|
||||
submit_button_component = gr.Button(
|
||||
value='Start ObjectClear', elem_id="start-button"
|
||||
)
|
||||
|
||||
with gr.Accordion('ObjectClear Settings', open=True):
|
||||
strength = gr.Radio(
|
||||
choices=[0.99, 1.0],
|
||||
value=0.99,
|
||||
label="Strength",
|
||||
info="0.99 better preserves the background and color; use 1.0 if object/shadow is not fully removed (default: 0.99)"
|
||||
)
|
||||
|
||||
guidance_scale = gr.Slider(
|
||||
minimum=1, maximum=10, step=0.5, value=2.5,
|
||||
label="Guidance Scale",
|
||||
info="Higher = stronger removal; lower = better background preservation (default: 2.5)"
|
||||
)
|
||||
|
||||
seed = gr.Slider(
|
||||
minimum=0, maximum=1000000, step=1, value=300000,
|
||||
label="Seed Value",
|
||||
info="Different seeds can lead to noticeably different object removal results (default: 300000)"
|
||||
)
|
||||
|
||||
num_inference_steps = gr.Slider(
|
||||
minimum=1, maximum=40, step=1, value=20,
|
||||
label="Num Inference Steps",
|
||||
info="Higher values may improve quality but take longer (default: 20)"
|
||||
)
|
||||
|
||||
|
||||
with gr.Column():
|
||||
output_image_component = gr.Image(
|
||||
type='pil', image_mode='RGB', label='Output', format="png", elem_id="input-image")
|
||||
|
||||
output_compare_image_component = gr.ImageSlider(
|
||||
label="Comparison",
|
||||
type="pil",
|
||||
format='png',
|
||||
elem_id="compare-image"
|
||||
)
|
||||
|
||||
input_image.upload(
|
||||
fn=upload_and_reset,
|
||||
inputs=[input_image, interactive_state],
|
||||
outputs=[
|
||||
image_state,
|
||||
image_info,
|
||||
input_image,
|
||||
interactive_state,
|
||||
click_state,
|
||||
mask_dropdown,
|
||||
]
|
||||
)
|
||||
|
||||
# click select image to get mask using sam
|
||||
input_image.select(
|
||||
fn=sam_refine,
|
||||
inputs=[image_state, point_prompt, click_state],
|
||||
outputs=[input_image, image_state, click_state]
|
||||
)
|
||||
|
||||
# add different mask
|
||||
add_mask_button.click(
|
||||
fn=add_multi_mask,
|
||||
inputs=[image_state, interactive_state, mask_dropdown],
|
||||
outputs=[interactive_state, mask_dropdown, input_image, click_state]
|
||||
)
|
||||
|
||||
remove_mask_button.click(
|
||||
fn=remove_multi_mask,
|
||||
inputs=[interactive_state, click_state, image_state],
|
||||
outputs=[interactive_state, mask_dropdown, input_image, click_state]
|
||||
)
|
||||
|
||||
# points clear
|
||||
clear_button_click.click(
|
||||
fn = clear_click,
|
||||
inputs = [image_state, click_state,],
|
||||
outputs = [input_image, click_state],
|
||||
)
|
||||
|
||||
submit_button_component.click(
|
||||
fn=process,
|
||||
inputs=[
|
||||
image_state,
|
||||
interactive_state,
|
||||
mask_dropdown,
|
||||
guidance_scale,
|
||||
seed,
|
||||
num_inference_steps,
|
||||
strength
|
||||
],
|
||||
outputs=[
|
||||
output_image_component, output_compare_image_component
|
||||
]
|
||||
)
|
||||
|
||||
with gr.Accordion("📕 Video Tutorial (click to expand)", open=False, elem_classes="custom-bg"):
|
||||
with gr.Row():
|
||||
gr.Video(value="./tutorial.mp4", elem_classes="video")
|
||||
|
||||
gr.Markdown("---")
|
||||
gr.Markdown("## Examples")
|
||||
|
||||
example_images = [
|
||||
os.path.join(os.path.dirname(__file__), "examples", f"test{i}.png")
|
||||
for i in range(10)
|
||||
]
|
||||
|
||||
examples_data = [
|
||||
[example_images[i], None] for i in range(len(example_images))
|
||||
]
|
||||
|
||||
examples = gr.Examples(
|
||||
examples=examples_data,
|
||||
inputs=[input_image, interactive_state],
|
||||
outputs=[image_state, image_info, input_image,
|
||||
interactive_state, click_state, mask_dropdown],
|
||||
fn=upload_and_reset,
|
||||
run_on_click=True,
|
||||
cache_examples=False,
|
||||
label="Click below to load example images"
|
||||
)
|
||||
|
||||
gr.Markdown(article)
|
||||
|
||||
|
||||
demo.launch(debug=True, show_error=True, share=True, server_port=args.port)
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 707 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 603 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 456 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 798 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 588 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 510 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 535 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 557 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 285 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 829 KiB |
@@ -0,0 +1,12 @@
|
||||
accelerate
|
||||
torch
|
||||
torchvision
|
||||
numpy
|
||||
opencv-python
|
||||
pillow
|
||||
transformers
|
||||
scipy
|
||||
diffusers
|
||||
segment-anything
|
||||
matplotlib
|
||||
gradio==5.30.0
|
||||
@@ -0,0 +1,129 @@
|
||||
import time
|
||||
import torch
|
||||
import cv2
|
||||
from PIL import Image, ImageDraw, ImageOps
|
||||
import numpy as np
|
||||
from typing import Union
|
||||
from segment_anything import sam_model_registry, SamPredictor, SamAutomaticMaskGenerator
|
||||
import matplotlib.pyplot as plt
|
||||
import PIL
|
||||
from .mask_painter import mask_painter
|
||||
|
||||
|
||||
class BaseSegmenter:
|
||||
def __init__(self, SAM_checkpoint, model_type, device='cuda:0'):
|
||||
"""
|
||||
device: model device
|
||||
SAM_checkpoint: path of SAM checkpoint
|
||||
model_type: vit_b, vit_l, vit_h
|
||||
"""
|
||||
print(f"Initializing BaseSegmenter to {device}")
|
||||
assert model_type in ['vit_b', 'vit_l', 'vit_h'], 'model_type must be vit_b, vit_l, or vit_h'
|
||||
|
||||
self.device = device
|
||||
self.torch_dtype = torch.float16 if 'cuda' in device else torch.float32
|
||||
self.model = sam_model_registry[model_type](checkpoint=SAM_checkpoint)
|
||||
self.model.to(device=self.device)
|
||||
self.predictor = SamPredictor(self.model)
|
||||
self.embedded = False
|
||||
|
||||
@torch.no_grad()
|
||||
def set_image(self, image: np.ndarray):
|
||||
# PIL.open(image_path) 3channel: RGB
|
||||
# image embedding: avoid encode the same image multiple times
|
||||
self.orignal_image = image
|
||||
if self.embedded:
|
||||
print('repeat embedding, please reset_image.')
|
||||
return
|
||||
self.predictor.set_image(image)
|
||||
self.embedded = True
|
||||
return
|
||||
|
||||
@torch.no_grad()
|
||||
def reset_image(self):
|
||||
# reset image embeding
|
||||
self.predictor.reset_image()
|
||||
self.embedded = False
|
||||
|
||||
def predict(self, prompts, mode, multimask=True):
|
||||
"""
|
||||
image: numpy array, h, w, 3
|
||||
prompts: dictionary, 3 keys: 'point_coords', 'point_labels', 'mask_input'
|
||||
prompts['point_coords']: numpy array [N,2]
|
||||
prompts['point_labels']: numpy array [1,N]
|
||||
prompts['mask_input']: numpy array [1,256,256]
|
||||
mode: 'point' (points only), 'mask' (mask only), 'both' (consider both)
|
||||
mask_outputs: True (return 3 masks), False (return 1 mask only)
|
||||
whem mask_outputs=True, mask_input=logits[np.argmax(scores), :, :][None, :, :]
|
||||
"""
|
||||
assert self.embedded, 'prediction is called before set_image (feature embedding).'
|
||||
assert mode in ['point', 'mask', 'both'], 'mode must be point, mask, or both'
|
||||
|
||||
if mode == 'point':
|
||||
masks, scores, logits = self.predictor.predict(point_coords=prompts['point_coords'],
|
||||
point_labels=prompts['point_labels'],
|
||||
multimask_output=multimask)
|
||||
elif mode == 'mask':
|
||||
masks, scores, logits = self.predictor.predict(mask_input=prompts['mask_input'],
|
||||
multimask_output=multimask)
|
||||
elif mode == 'both': # both
|
||||
masks, scores, logits = self.predictor.predict(point_coords=prompts['point_coords'],
|
||||
point_labels=prompts['point_labels'],
|
||||
mask_input=prompts['mask_input'],
|
||||
multimask_output=multimask)
|
||||
else:
|
||||
raise("Not implement now!")
|
||||
# masks (n, h, w), scores (n,), logits (n, 256, 256)
|
||||
return masks, scores, logits
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# load and show an image
|
||||
image = cv2.imread('/hhd3/gaoshang/truck.jpg')
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # numpy array (h, w, 3)
|
||||
|
||||
# initialise BaseSegmenter
|
||||
SAM_checkpoint= '/ssd1/gaomingqi/checkpoints/sam_vit_h_4b8939.pth'
|
||||
model_type = 'vit_h'
|
||||
device = "cuda:4"
|
||||
base_segmenter = BaseSegmenter(SAM_checkpoint=SAM_checkpoint, model_type=model_type, device=device)
|
||||
|
||||
# image embedding (once embedded, multiple prompts can be applied)
|
||||
base_segmenter.set_image(image)
|
||||
|
||||
# examples
|
||||
# point only ------------------------
|
||||
mode = 'point'
|
||||
prompts = {
|
||||
'point_coords': np.array([[500, 375], [1125, 625]]),
|
||||
'point_labels': np.array([1, 1]),
|
||||
}
|
||||
masks, scores, logits = base_segmenter.predict(prompts, mode, multimask=False) # masks (n, h, w), scores (n,), logits (n, 256, 256)
|
||||
painted_image = mask_painter(image, masks[np.argmax(scores)].astype('uint8'), background_alpha=0.8)
|
||||
painted_image = cv2.cvtColor(painted_image, cv2.COLOR_RGB2BGR) # numpy array (h, w, 3)
|
||||
cv2.imwrite('/hhd3/gaoshang/truck_point.jpg', painted_image)
|
||||
|
||||
# both ------------------------
|
||||
mode = 'both'
|
||||
mask_input = logits[np.argmax(scores), :, :]
|
||||
prompts = {'mask_input': mask_input [None, :, :]}
|
||||
prompts = {
|
||||
'point_coords': np.array([[500, 375], [1125, 625]]),
|
||||
'point_labels': np.array([1, 0]),
|
||||
'mask_input': mask_input[None, :, :]
|
||||
}
|
||||
masks, scores, logits = base_segmenter.predict(prompts, mode, multimask=True) # masks (n, h, w), scores (n,), logits (n, 256, 256)
|
||||
painted_image = mask_painter(image, masks[np.argmax(scores)].astype('uint8'), background_alpha=0.8)
|
||||
painted_image = cv2.cvtColor(painted_image, cv2.COLOR_RGB2BGR) # numpy array (h, w, 3)
|
||||
cv2.imwrite('/hhd3/gaoshang/truck_both.jpg', painted_image)
|
||||
|
||||
# mask only ------------------------
|
||||
mode = 'mask'
|
||||
mask_input = logits[np.argmax(scores), :, :]
|
||||
|
||||
prompts = {'mask_input': mask_input[None, :, :]}
|
||||
|
||||
masks, scores, logits = base_segmenter.predict(prompts, mode, multimask=True) # masks (n, h, w), scores (n,), logits (n, 256, 256)
|
||||
painted_image = mask_painter(image, masks[np.argmax(scores)].astype('uint8'), background_alpha=0.8)
|
||||
painted_image = cv2.cvtColor(painted_image, cv2.COLOR_RGB2BGR) # numpy array (h, w, 3)
|
||||
cv2.imwrite('/hhd3/gaoshang/truck_mask.jpg', painted_image)
|
||||
@@ -0,0 +1,109 @@
|
||||
import math
|
||||
import os
|
||||
import requests
|
||||
from torch.hub import download_url_to_file, get_dir
|
||||
from tqdm import tqdm
|
||||
from urllib.parse import urlparse
|
||||
|
||||
def sizeof_fmt(size, suffix='B'):
|
||||
"""Get human readable file size.
|
||||
|
||||
Args:
|
||||
size (int): File size.
|
||||
suffix (str): Suffix. Default: 'B'.
|
||||
|
||||
Return:
|
||||
str: Formated file siz.
|
||||
"""
|
||||
for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']:
|
||||
if abs(size) < 1024.0:
|
||||
return f'{size:3.1f} {unit}{suffix}'
|
||||
size /= 1024.0
|
||||
return f'{size:3.1f} Y{suffix}'
|
||||
|
||||
|
||||
def download_file_from_google_drive(file_id, save_path):
|
||||
"""Download files from google drive.
|
||||
Ref:
|
||||
https://stackoverflow.com/questions/25010369/wget-curl-large-file-from-google-drive # noqa E501
|
||||
Args:
|
||||
file_id (str): File id.
|
||||
save_path (str): Save path.
|
||||
"""
|
||||
|
||||
session = requests.Session()
|
||||
URL = 'https://docs.google.com/uc?export=download'
|
||||
params = {'id': file_id}
|
||||
|
||||
response = session.get(URL, params=params, stream=True)
|
||||
token = get_confirm_token(response)
|
||||
if token:
|
||||
params['confirm'] = token
|
||||
response = session.get(URL, params=params, stream=True)
|
||||
|
||||
# get file size
|
||||
response_file_size = session.get(URL, params=params, stream=True, headers={'Range': 'bytes=0-2'})
|
||||
print(response_file_size)
|
||||
if 'Content-Range' in response_file_size.headers:
|
||||
file_size = int(response_file_size.headers['Content-Range'].split('/')[1])
|
||||
else:
|
||||
file_size = None
|
||||
|
||||
save_response_content(response, save_path, file_size)
|
||||
|
||||
|
||||
def get_confirm_token(response):
|
||||
for key, value in response.cookies.items():
|
||||
if key.startswith('download_warning'):
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
def save_response_content(response, destination, file_size=None, chunk_size=32768):
|
||||
if file_size is not None:
|
||||
pbar = tqdm(total=math.ceil(file_size / chunk_size), unit='chunk')
|
||||
|
||||
readable_file_size = sizeof_fmt(file_size)
|
||||
else:
|
||||
pbar = None
|
||||
|
||||
with open(destination, 'wb') as f:
|
||||
downloaded_size = 0
|
||||
for chunk in response.iter_content(chunk_size):
|
||||
downloaded_size += chunk_size
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
pbar.set_description(f'Download {sizeof_fmt(downloaded_size)} / {readable_file_size}')
|
||||
if chunk: # filter out keep-alive new chunks
|
||||
f.write(chunk)
|
||||
if pbar is not None:
|
||||
pbar.close()
|
||||
|
||||
|
||||
def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
|
||||
"""Load file form http url, will download models if necessary.
|
||||
Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py
|
||||
Args:
|
||||
url (str): URL to be downloaded.
|
||||
model_dir (str): The path to save the downloaded model. Should be a full path. If None, use pytorch hub_dir.
|
||||
Default: None.
|
||||
progress (bool): Whether to show the download progress. Default: True.
|
||||
file_name (str): The downloaded file name. If None, use the file name in the url. Default: None.
|
||||
Returns:
|
||||
str: The path to the downloaded file.
|
||||
"""
|
||||
if model_dir is None: # use the pytorch hub_dir
|
||||
hub_dir = get_dir()
|
||||
model_dir = os.path.join(hub_dir, 'checkpoints')
|
||||
|
||||
os.makedirs(model_dir, exist_ok=True)
|
||||
|
||||
parts = urlparse(url)
|
||||
filename = os.path.basename(parts.path)
|
||||
if file_name is not None:
|
||||
filename = file_name
|
||||
cached_file = os.path.abspath(os.path.join(model_dir, filename))
|
||||
if not os.path.exists(cached_file):
|
||||
print(f'Downloading: "{url}" to {cached_file}\n')
|
||||
download_url_to_file(url, cached_file, hash_prefix=None, progress=progress)
|
||||
return cached_file
|
||||
@@ -0,0 +1,99 @@
|
||||
import time
|
||||
import torch
|
||||
import cv2
|
||||
from PIL import Image, ImageDraw, ImageOps
|
||||
import numpy as np
|
||||
from typing import Union
|
||||
from segment_anything import sam_model_registry, SamPredictor, SamAutomaticMaskGenerator
|
||||
import matplotlib.pyplot as plt
|
||||
import PIL
|
||||
from .mask_painter import mask_painter as mask_painter2
|
||||
from .base_segmenter import BaseSegmenter
|
||||
from .painter import mask_painter, point_painter
|
||||
import os
|
||||
import requests
|
||||
import sys
|
||||
|
||||
|
||||
mask_color = 3
|
||||
mask_alpha = 0.7
|
||||
contour_color = 1
|
||||
contour_width = 5
|
||||
point_color_ne = 8
|
||||
point_color_ps = 50
|
||||
point_alpha = 0.9
|
||||
point_radius = 15
|
||||
contour_color = 2
|
||||
contour_width = 5
|
||||
|
||||
|
||||
class SamControler():
|
||||
def __init__(self, SAM_checkpoint, model_type, device):
|
||||
'''
|
||||
initialize sam controler
|
||||
'''
|
||||
self.sam_controler = BaseSegmenter(SAM_checkpoint, model_type, device)
|
||||
|
||||
|
||||
# def seg_again(self, image: np.ndarray):
|
||||
# '''
|
||||
# it is used when interact in video
|
||||
# '''
|
||||
# self.sam_controler.reset_image()
|
||||
# self.sam_controler.set_image(image)
|
||||
# return
|
||||
|
||||
|
||||
def first_frame_click(self, image: np.ndarray, points:np.ndarray, labels: np.ndarray, multimask=True,mask_color=3):
|
||||
'''
|
||||
it is used in first frame in video
|
||||
return: mask, logit, painted image(mask+point)
|
||||
'''
|
||||
# self.sam_controler.set_image(image)
|
||||
origal_image = self.sam_controler.orignal_image
|
||||
neg_flag = labels[-1]
|
||||
if neg_flag==1:
|
||||
#find neg
|
||||
prompts = {
|
||||
'point_coords': points,
|
||||
'point_labels': labels,
|
||||
}
|
||||
masks, scores, logits = self.sam_controler.predict(prompts, 'point', multimask)
|
||||
mask, logit = masks[np.argmax(scores)], logits[np.argmax(scores), :, :]
|
||||
prompts = {
|
||||
'point_coords': points,
|
||||
'point_labels': labels,
|
||||
'mask_input': logit[None, :, :]
|
||||
}
|
||||
masks, scores, logits = self.sam_controler.predict(prompts, 'both', multimask)
|
||||
mask, logit = masks[np.argmax(scores)], logits[np.argmax(scores), :, :]
|
||||
else:
|
||||
#find positive
|
||||
prompts = {
|
||||
'point_coords': points,
|
||||
'point_labels': labels,
|
||||
}
|
||||
masks, scores, logits = self.sam_controler.predict(prompts, 'point', multimask)
|
||||
mask, logit = masks[np.argmax(scores)], logits[np.argmax(scores), :, :]
|
||||
|
||||
|
||||
assert len(points)==len(labels)
|
||||
|
||||
painted_image = mask_painter(image, mask.astype('uint8'), mask_color, mask_alpha, contour_color, contour_width)
|
||||
painted_image = point_painter(painted_image, np.squeeze(points[np.argwhere(labels>0)],axis = 1), point_color_ne, point_alpha, point_radius, contour_color, contour_width)
|
||||
painted_image = point_painter(painted_image, np.squeeze(points[np.argwhere(labels<1)],axis = 1), point_color_ps, point_alpha, point_radius, contour_color, contour_width)
|
||||
painted_image = Image.fromarray(painted_image)
|
||||
|
||||
return mask, logit, painted_image
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,288 @@
|
||||
import cv2
|
||||
import torch
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import copy
|
||||
import time
|
||||
|
||||
|
||||
def colormap(rgb=True):
|
||||
color_list = np.array(
|
||||
[
|
||||
0.000, 0.000, 0.000,
|
||||
1.000, 1.000, 1.000,
|
||||
1.000, 0.498, 0.313,
|
||||
0.392, 0.581, 0.929,
|
||||
0.000, 0.447, 0.741,
|
||||
0.850, 0.325, 0.098,
|
||||
0.929, 0.694, 0.125,
|
||||
0.494, 0.184, 0.556,
|
||||
0.466, 0.674, 0.188,
|
||||
0.301, 0.745, 0.933,
|
||||
0.635, 0.078, 0.184,
|
||||
0.300, 0.300, 0.300,
|
||||
0.600, 0.600, 0.600,
|
||||
1.000, 0.000, 0.000,
|
||||
1.000, 0.500, 0.000,
|
||||
0.749, 0.749, 0.000,
|
||||
0.000, 1.000, 0.000,
|
||||
0.000, 0.000, 1.000,
|
||||
0.667, 0.000, 1.000,
|
||||
0.333, 0.333, 0.000,
|
||||
0.333, 0.667, 0.000,
|
||||
0.333, 1.000, 0.000,
|
||||
0.667, 0.333, 0.000,
|
||||
0.667, 0.667, 0.000,
|
||||
0.667, 1.000, 0.000,
|
||||
1.000, 0.333, 0.000,
|
||||
1.000, 0.667, 0.000,
|
||||
1.000, 1.000, 0.000,
|
||||
0.000, 0.333, 0.500,
|
||||
0.000, 0.667, 0.500,
|
||||
0.000, 1.000, 0.500,
|
||||
0.333, 0.000, 0.500,
|
||||
0.333, 0.333, 0.500,
|
||||
0.333, 0.667, 0.500,
|
||||
0.333, 1.000, 0.500,
|
||||
0.667, 0.000, 0.500,
|
||||
0.667, 0.333, 0.500,
|
||||
0.667, 0.667, 0.500,
|
||||
0.667, 1.000, 0.500,
|
||||
1.000, 0.000, 0.500,
|
||||
1.000, 0.333, 0.500,
|
||||
1.000, 0.667, 0.500,
|
||||
1.000, 1.000, 0.500,
|
||||
0.000, 0.333, 1.000,
|
||||
0.000, 0.667, 1.000,
|
||||
0.000, 1.000, 1.000,
|
||||
0.333, 0.000, 1.000,
|
||||
0.333, 0.333, 1.000,
|
||||
0.333, 0.667, 1.000,
|
||||
0.333, 1.000, 1.000,
|
||||
0.667, 0.000, 1.000,
|
||||
0.667, 0.333, 1.000,
|
||||
0.667, 0.667, 1.000,
|
||||
0.667, 1.000, 1.000,
|
||||
1.000, 0.000, 1.000,
|
||||
1.000, 0.333, 1.000,
|
||||
1.000, 0.667, 1.000,
|
||||
0.167, 0.000, 0.000,
|
||||
0.333, 0.000, 0.000,
|
||||
0.500, 0.000, 0.000,
|
||||
0.667, 0.000, 0.000,
|
||||
0.833, 0.000, 0.000,
|
||||
1.000, 0.000, 0.000,
|
||||
0.000, 0.167, 0.000,
|
||||
0.000, 0.333, 0.000,
|
||||
0.000, 0.500, 0.000,
|
||||
0.000, 0.667, 0.000,
|
||||
0.000, 0.833, 0.000,
|
||||
0.000, 1.000, 0.000,
|
||||
0.000, 0.000, 0.167,
|
||||
0.000, 0.000, 0.333,
|
||||
0.000, 0.000, 0.500,
|
||||
0.000, 0.000, 0.667,
|
||||
0.000, 0.000, 0.833,
|
||||
0.000, 0.000, 1.000,
|
||||
0.143, 0.143, 0.143,
|
||||
0.286, 0.286, 0.286,
|
||||
0.429, 0.429, 0.429,
|
||||
0.571, 0.571, 0.571,
|
||||
0.714, 0.714, 0.714,
|
||||
0.857, 0.857, 0.857
|
||||
]
|
||||
).astype(np.float32)
|
||||
color_list = color_list.reshape((-1, 3)) * 255
|
||||
if not rgb:
|
||||
color_list = color_list[:, ::-1]
|
||||
return color_list
|
||||
|
||||
|
||||
color_list = colormap()
|
||||
color_list = color_list.astype('uint8').tolist()
|
||||
|
||||
|
||||
def vis_add_mask(image, background_mask, contour_mask, background_color, contour_color, background_alpha, contour_alpha):
|
||||
background_color = np.array(background_color)
|
||||
contour_color = np.array(contour_color)
|
||||
|
||||
# background_mask = 1 - background_mask
|
||||
# contour_mask = 1 - contour_mask
|
||||
|
||||
for i in range(3):
|
||||
image[:, :, i] = image[:, :, i] * (1-background_alpha+background_mask*background_alpha) \
|
||||
+ background_color[i] * (background_alpha-background_mask*background_alpha)
|
||||
|
||||
image[:, :, i] = image[:, :, i] * (1-contour_alpha+contour_mask*contour_alpha) \
|
||||
+ contour_color[i] * (contour_alpha-contour_mask*contour_alpha)
|
||||
|
||||
return image.astype('uint8')
|
||||
|
||||
|
||||
def mask_generator_00(mask, background_radius, contour_radius):
|
||||
# no background width when '00'
|
||||
# distance map
|
||||
dist_transform_fore = cv2.distanceTransform(mask, cv2.DIST_L2, 3)
|
||||
dist_transform_back = cv2.distanceTransform(1-mask, cv2.DIST_L2, 3)
|
||||
dist_map = dist_transform_fore - dist_transform_back
|
||||
# ...:::!!!:::...
|
||||
contour_radius += 2
|
||||
contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius))
|
||||
contour_mask = contour_mask / np.max(contour_mask)
|
||||
contour_mask[contour_mask>0.5] = 1.
|
||||
|
||||
return mask, contour_mask
|
||||
|
||||
|
||||
def mask_generator_01(mask, background_radius, contour_radius):
|
||||
# no background width when '00'
|
||||
# distance map
|
||||
dist_transform_fore = cv2.distanceTransform(mask, cv2.DIST_L2, 3)
|
||||
dist_transform_back = cv2.distanceTransform(1-mask, cv2.DIST_L2, 3)
|
||||
dist_map = dist_transform_fore - dist_transform_back
|
||||
# ...:::!!!:::...
|
||||
contour_radius += 2
|
||||
contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius))
|
||||
contour_mask = contour_mask / np.max(contour_mask)
|
||||
return mask, contour_mask
|
||||
|
||||
|
||||
def mask_generator_10(mask, background_radius, contour_radius):
|
||||
# distance map
|
||||
dist_transform_fore = cv2.distanceTransform(mask, cv2.DIST_L2, 3)
|
||||
dist_transform_back = cv2.distanceTransform(1-mask, cv2.DIST_L2, 3)
|
||||
dist_map = dist_transform_fore - dist_transform_back
|
||||
# .....:::::!!!!!
|
||||
background_mask = np.clip(dist_map, -background_radius, background_radius)
|
||||
background_mask = (background_mask - np.min(background_mask))
|
||||
background_mask = background_mask / np.max(background_mask)
|
||||
# ...:::!!!:::...
|
||||
contour_radius += 2
|
||||
contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius))
|
||||
contour_mask = contour_mask / np.max(contour_mask)
|
||||
contour_mask[contour_mask>0.5] = 1.
|
||||
return background_mask, contour_mask
|
||||
|
||||
|
||||
def mask_generator_11(mask, background_radius, contour_radius):
|
||||
# distance map
|
||||
dist_transform_fore = cv2.distanceTransform(mask, cv2.DIST_L2, 3)
|
||||
dist_transform_back = cv2.distanceTransform(1-mask, cv2.DIST_L2, 3)
|
||||
dist_map = dist_transform_fore - dist_transform_back
|
||||
# .....:::::!!!!!
|
||||
background_mask = np.clip(dist_map, -background_radius, background_radius)
|
||||
background_mask = (background_mask - np.min(background_mask))
|
||||
background_mask = background_mask / np.max(background_mask)
|
||||
# ...:::!!!:::...
|
||||
contour_radius += 2
|
||||
contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius))
|
||||
contour_mask = contour_mask / np.max(contour_mask)
|
||||
return background_mask, contour_mask
|
||||
|
||||
|
||||
def mask_painter(input_image, input_mask, background_alpha=0.5, background_blur_radius=7, contour_width=3, contour_color=3, contour_alpha=1, mode='11'):
|
||||
"""
|
||||
Input:
|
||||
input_image: numpy array
|
||||
input_mask: numpy array
|
||||
background_alpha: transparency of background, [0, 1], 1: all black, 0: do nothing
|
||||
background_blur_radius: radius of background blur, must be odd number
|
||||
contour_width: width of mask contour, must be odd number
|
||||
contour_color: color index (in color map) of mask contour, 0: black, 1: white, >1: others
|
||||
contour_alpha: transparency of mask contour, [0, 1], if 0: no contour highlighted
|
||||
mode: painting mode, '00', no blur, '01' only blur contour, '10' only blur background, '11' blur both
|
||||
|
||||
Output:
|
||||
painted_image: numpy array
|
||||
"""
|
||||
assert input_image.shape[:2] == input_mask.shape, 'different shape'
|
||||
assert background_blur_radius % 2 * contour_width % 2 > 0, 'background_blur_radius and contour_width must be ODD'
|
||||
assert mode in ['00', '01', '10', '11'], 'mode should be 00, 01, 10, or 11'
|
||||
|
||||
# downsample input image and mask
|
||||
width, height = input_image.shape[0], input_image.shape[1]
|
||||
res = 1024
|
||||
ratio = min(1.0 * res / max(width, height), 1.0)
|
||||
input_image = cv2.resize(input_image, (int(height*ratio), int(width*ratio)))
|
||||
input_mask = cv2.resize(input_mask, (int(height*ratio), int(width*ratio)))
|
||||
|
||||
# 0: background, 1: foreground
|
||||
msk = np.clip(input_mask, 0, 1)
|
||||
|
||||
# generate masks for background and contour pixels
|
||||
background_radius = (background_blur_radius - 1) // 2
|
||||
contour_radius = (contour_width - 1) // 2
|
||||
generator_dict = {'00':mask_generator_00, '01':mask_generator_01, '10':mask_generator_10, '11':mask_generator_11}
|
||||
background_mask, contour_mask = generator_dict[mode](msk, background_radius, contour_radius)
|
||||
|
||||
# paint
|
||||
painted_image = vis_add_mask\
|
||||
(input_image, background_mask, contour_mask, color_list[0], color_list[contour_color], background_alpha, contour_alpha) # black for background
|
||||
|
||||
return painted_image
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
background_alpha = 0.7 # transparency of background 1: all black, 0: do nothing
|
||||
background_blur_radius = 31 # radius of background blur, must be odd number
|
||||
contour_width = 11 # contour width, must be odd number
|
||||
contour_color = 3 # id in color map, 0: black, 1: white, >1: others
|
||||
contour_alpha = 1 # transparency of background, 0: no contour highlighted
|
||||
|
||||
# load input image and mask
|
||||
input_image = np.array(Image.open('./test_img/painter_input_image.jpg').convert('RGB'))
|
||||
input_mask = np.array(Image.open('./test_img/painter_input_mask.jpg').convert('P'))
|
||||
|
||||
# paint
|
||||
overall_time_1 = 0
|
||||
overall_time_2 = 0
|
||||
overall_time_3 = 0
|
||||
overall_time_4 = 0
|
||||
overall_time_5 = 0
|
||||
|
||||
for i in range(50):
|
||||
t2 = time.time()
|
||||
painted_image_00 = mask_painter(input_image, input_mask, background_alpha, background_blur_radius, contour_width, contour_color, contour_alpha, mode='00')
|
||||
e2 = time.time()
|
||||
|
||||
t3 = time.time()
|
||||
painted_image_10 = mask_painter(input_image, input_mask, background_alpha, background_blur_radius, contour_width, contour_color, contour_alpha, mode='10')
|
||||
e3 = time.time()
|
||||
|
||||
t1 = time.time()
|
||||
painted_image = mask_painter(input_image, input_mask, background_alpha, background_blur_radius, contour_width, contour_color, contour_alpha)
|
||||
e1 = time.time()
|
||||
|
||||
t4 = time.time()
|
||||
painted_image_01 = mask_painter(input_image, input_mask, background_alpha, background_blur_radius, contour_width, contour_color, contour_alpha, mode='01')
|
||||
e4 = time.time()
|
||||
|
||||
t5 = time.time()
|
||||
painted_image_11 = mask_painter(input_image, input_mask, background_alpha, background_blur_radius, contour_width, contour_color, contour_alpha, mode='11')
|
||||
e5 = time.time()
|
||||
|
||||
overall_time_1 += (e1 - t1)
|
||||
overall_time_2 += (e2 - t2)
|
||||
overall_time_3 += (e3 - t3)
|
||||
overall_time_4 += (e4 - t4)
|
||||
overall_time_5 += (e5 - t5)
|
||||
|
||||
print(f'average time w gaussian: {overall_time_1/50}')
|
||||
print(f'average time w/o gaussian00: {overall_time_2/50}')
|
||||
print(f'average time w/o gaussian10: {overall_time_3/50}')
|
||||
print(f'average time w/o gaussian01: {overall_time_4/50}')
|
||||
print(f'average time w/o gaussian11: {overall_time_5/50}')
|
||||
|
||||
# save
|
||||
painted_image_00 = Image.fromarray(painted_image_00)
|
||||
painted_image_00.save('./test_img/painter_output_image_00.png')
|
||||
|
||||
painted_image_10 = Image.fromarray(painted_image_10)
|
||||
painted_image_10.save('./test_img/painter_output_image_10.png')
|
||||
|
||||
painted_image_01 = Image.fromarray(painted_image_01)
|
||||
painted_image_01.save('./test_img/painter_output_image_01.png')
|
||||
|
||||
painted_image_11 = Image.fromarray(painted_image_11)
|
||||
painted_image_11.save('./test_img/painter_output_image_11.png')
|
||||
@@ -0,0 +1,131 @@
|
||||
import os
|
||||
import re
|
||||
import random
|
||||
import time
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import logging
|
||||
import numpy as np
|
||||
from os import path as osp
|
||||
|
||||
def constant_init(module, val, bias=0):
|
||||
if hasattr(module, 'weight') and module.weight is not None:
|
||||
nn.init.constant_(module.weight, val)
|
||||
if hasattr(module, 'bias') and module.bias is not None:
|
||||
nn.init.constant_(module.bias, bias)
|
||||
|
||||
initialized_logger = {}
|
||||
def get_root_logger(logger_name='basicsr', log_level=logging.INFO, log_file=None):
|
||||
"""Get the root logger.
|
||||
The logger will be initialized if it has not been initialized. By default a
|
||||
StreamHandler will be added. If `log_file` is specified, a FileHandler will
|
||||
also be added.
|
||||
Args:
|
||||
logger_name (str): root logger name. Default: 'basicsr'.
|
||||
log_file (str | None): The log filename. If specified, a FileHandler
|
||||
will be added to the root logger.
|
||||
log_level (int): The root logger level. Note that only the process of
|
||||
rank 0 is affected, while other processes will set the level to
|
||||
"Error" and be silent most of the time.
|
||||
Returns:
|
||||
logging.Logger: The root logger.
|
||||
"""
|
||||
logger = logging.getLogger(logger_name)
|
||||
# if the logger has been initialized, just return it
|
||||
if logger_name in initialized_logger:
|
||||
return logger
|
||||
|
||||
format_str = '%(asctime)s %(levelname)s: %(message)s'
|
||||
stream_handler = logging.StreamHandler()
|
||||
stream_handler.setFormatter(logging.Formatter(format_str))
|
||||
logger.addHandler(stream_handler)
|
||||
logger.propagate = False
|
||||
|
||||
if log_file is not None:
|
||||
logger.setLevel(log_level)
|
||||
# add file handler
|
||||
# file_handler = logging.FileHandler(log_file, 'w')
|
||||
file_handler = logging.FileHandler(log_file, 'a') #Shangchen: keep the previous log
|
||||
file_handler.setFormatter(logging.Formatter(format_str))
|
||||
file_handler.setLevel(log_level)
|
||||
logger.addHandler(file_handler)
|
||||
initialized_logger[logger_name] = True
|
||||
return logger
|
||||
|
||||
|
||||
IS_HIGH_VERSION = [int(m) for m in list(re.findall(r"^([0-9]+)\.([0-9]+)\.([0-9]+)([^0-9][a-zA-Z0-9]*)?(\+git.*)?$",\
|
||||
torch.__version__)[0][:3])] >= [1, 12, 0]
|
||||
|
||||
def gpu_is_available():
|
||||
if IS_HIGH_VERSION:
|
||||
if torch.backends.mps.is_available():
|
||||
return True
|
||||
return True if torch.cuda.is_available() and torch.backends.cudnn.is_available() else False
|
||||
|
||||
def get_device(gpu_id=None):
|
||||
if gpu_id is None:
|
||||
gpu_str = ''
|
||||
elif isinstance(gpu_id, int):
|
||||
gpu_str = f':{gpu_id}'
|
||||
else:
|
||||
raise TypeError('Input should be int value.')
|
||||
|
||||
if IS_HIGH_VERSION:
|
||||
if torch.backends.mps.is_available():
|
||||
return torch.device('mps'+gpu_str)
|
||||
return torch.device('cuda'+gpu_str if torch.cuda.is_available() and torch.backends.cudnn.is_available() else 'cpu')
|
||||
|
||||
|
||||
def set_random_seed(seed):
|
||||
"""Set random seeds."""
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
|
||||
def get_time_str():
|
||||
return time.strftime('%Y%m%d_%H%M%S', time.localtime())
|
||||
|
||||
|
||||
def scandir(dir_path, suffix=None, recursive=False, full_path=False):
|
||||
"""Scan a directory to find the interested files.
|
||||
|
||||
Args:
|
||||
dir_path (str): Path of the directory.
|
||||
suffix (str | tuple(str), optional): File suffix that we are
|
||||
interested in. Default: None.
|
||||
recursive (bool, optional): If set to True, recursively scan the
|
||||
directory. Default: False.
|
||||
full_path (bool, optional): If set to True, include the dir_path.
|
||||
Default: False.
|
||||
|
||||
Returns:
|
||||
A generator for all the interested files with relative pathes.
|
||||
"""
|
||||
|
||||
if (suffix is not None) and not isinstance(suffix, (str, tuple)):
|
||||
raise TypeError('"suffix" must be a string or tuple of strings')
|
||||
|
||||
root = dir_path
|
||||
|
||||
def _scandir(dir_path, suffix, recursive):
|
||||
for entry in os.scandir(dir_path):
|
||||
if not entry.name.startswith('.') and entry.is_file():
|
||||
if full_path:
|
||||
return_path = entry.path
|
||||
else:
|
||||
return_path = osp.relpath(entry.path, root)
|
||||
|
||||
if suffix is None:
|
||||
yield return_path
|
||||
elif return_path.endswith(suffix):
|
||||
yield return_path
|
||||
else:
|
||||
if recursive:
|
||||
yield from _scandir(entry.path, suffix=suffix, recursive=recursive)
|
||||
else:
|
||||
continue
|
||||
|
||||
return _scandir(dir_path, suffix=suffix, recursive=recursive)
|
||||
@@ -0,0 +1,215 @@
|
||||
# paint masks, contours, or points on images, with specified colors
|
||||
import cv2
|
||||
import torch
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import copy
|
||||
import time
|
||||
|
||||
|
||||
def colormap(rgb=True):
|
||||
color_list = np.array(
|
||||
[
|
||||
0.000, 0.000, 0.000,
|
||||
1.000, 1.000, 1.000,
|
||||
1.000, 0.498, 0.313,
|
||||
0.392, 0.581, 0.929,
|
||||
0.000, 0.447, 0.741,
|
||||
0.850, 0.325, 0.098,
|
||||
0.929, 0.694, 0.125,
|
||||
0.494, 0.184, 0.556,
|
||||
0.466, 0.674, 0.188,
|
||||
0.301, 0.745, 0.933,
|
||||
0.635, 0.078, 0.184,
|
||||
0.300, 0.300, 0.300,
|
||||
0.600, 0.600, 0.600,
|
||||
1.000, 0.000, 0.000,
|
||||
1.000, 0.500, 0.000,
|
||||
0.749, 0.749, 0.000,
|
||||
0.000, 1.000, 0.000,
|
||||
0.000, 0.000, 1.000,
|
||||
0.667, 0.000, 1.000,
|
||||
0.333, 0.333, 0.000,
|
||||
0.333, 0.667, 0.000,
|
||||
0.333, 1.000, 0.000,
|
||||
0.667, 0.333, 0.000,
|
||||
0.667, 0.667, 0.000,
|
||||
0.667, 1.000, 0.000,
|
||||
1.000, 0.333, 0.000,
|
||||
1.000, 0.667, 0.000,
|
||||
1.000, 1.000, 0.000,
|
||||
0.000, 0.333, 0.500,
|
||||
0.000, 0.667, 0.500,
|
||||
0.000, 1.000, 0.500,
|
||||
0.333, 0.000, 0.500,
|
||||
0.333, 0.333, 0.500,
|
||||
0.333, 0.667, 0.500,
|
||||
0.333, 1.000, 0.500,
|
||||
0.667, 0.000, 0.500,
|
||||
0.667, 0.333, 0.500,
|
||||
0.667, 0.667, 0.500,
|
||||
0.667, 1.000, 0.500,
|
||||
1.000, 0.000, 0.500,
|
||||
1.000, 0.333, 0.500,
|
||||
1.000, 0.667, 0.500,
|
||||
1.000, 1.000, 0.500,
|
||||
0.000, 0.333, 1.000,
|
||||
0.000, 0.667, 1.000,
|
||||
0.000, 1.000, 1.000,
|
||||
0.333, 0.000, 1.000,
|
||||
0.333, 0.333, 1.000,
|
||||
0.333, 0.667, 1.000,
|
||||
0.333, 1.000, 1.000,
|
||||
0.667, 0.000, 1.000,
|
||||
0.667, 0.333, 1.000,
|
||||
0.667, 0.667, 1.000,
|
||||
0.667, 1.000, 1.000,
|
||||
1.000, 0.000, 1.000,
|
||||
1.000, 0.333, 1.000,
|
||||
1.000, 0.667, 1.000,
|
||||
0.167, 0.000, 0.000,
|
||||
0.333, 0.000, 0.000,
|
||||
0.500, 0.000, 0.000,
|
||||
0.667, 0.000, 0.000,
|
||||
0.833, 0.000, 0.000,
|
||||
1.000, 0.000, 0.000,
|
||||
0.000, 0.167, 0.000,
|
||||
0.000, 0.333, 0.000,
|
||||
0.000, 0.500, 0.000,
|
||||
0.000, 0.667, 0.000,
|
||||
0.000, 0.833, 0.000,
|
||||
0.000, 1.000, 0.000,
|
||||
0.000, 0.000, 0.167,
|
||||
0.000, 0.000, 0.333,
|
||||
0.000, 0.000, 0.500,
|
||||
0.000, 0.000, 0.667,
|
||||
0.000, 0.000, 0.833,
|
||||
0.000, 0.000, 1.000,
|
||||
0.143, 0.143, 0.143,
|
||||
0.286, 0.286, 0.286,
|
||||
0.429, 0.429, 0.429,
|
||||
0.571, 0.571, 0.571,
|
||||
0.714, 0.714, 0.714,
|
||||
0.857, 0.857, 0.857
|
||||
]
|
||||
).astype(np.float32)
|
||||
color_list = color_list.reshape((-1, 3)) * 255
|
||||
if not rgb:
|
||||
color_list = color_list[:, ::-1]
|
||||
return color_list
|
||||
|
||||
|
||||
color_list = colormap()
|
||||
color_list = color_list.astype('uint8').tolist()
|
||||
|
||||
|
||||
def vis_add_mask(image, mask, color, alpha):
|
||||
color = np.array(color_list[color])
|
||||
mask = mask > 0.5
|
||||
image[mask] = image[mask] * (1-alpha) + color * alpha
|
||||
return image.astype('uint8')
|
||||
|
||||
def point_painter(input_image, input_points, point_color=5, point_alpha=0.9, point_radius=15, contour_color=2, contour_width=5):
|
||||
h, w = input_image.shape[:2]
|
||||
point_mask = np.zeros((h, w)).astype('uint8')
|
||||
for point in input_points:
|
||||
point_mask[point[1], point[0]] = 1
|
||||
|
||||
kernel = cv2.getStructuringElement(2, (point_radius, point_radius))
|
||||
point_mask = cv2.dilate(point_mask, kernel)
|
||||
|
||||
contour_radius = (contour_width - 1) // 2
|
||||
dist_transform_fore = cv2.distanceTransform(point_mask, cv2.DIST_L2, 3)
|
||||
dist_transform_back = cv2.distanceTransform(1-point_mask, cv2.DIST_L2, 3)
|
||||
dist_map = dist_transform_fore - dist_transform_back
|
||||
# ...:::!!!:::...
|
||||
contour_radius += 2
|
||||
contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius))
|
||||
contour_mask = contour_mask / np.max(contour_mask)
|
||||
contour_mask[contour_mask>0.5] = 1.
|
||||
|
||||
# paint mask
|
||||
painted_image = vis_add_mask(input_image.copy(), point_mask, point_color, point_alpha)
|
||||
# paint contour
|
||||
painted_image = vis_add_mask(painted_image.copy(), 1-contour_mask, contour_color, 1)
|
||||
return painted_image
|
||||
|
||||
def mask_painter(input_image, input_mask, mask_color=5, mask_alpha=0.7, contour_color=1, contour_width=3):
|
||||
assert input_image.shape[:2] == input_mask.shape, 'different shape between image and mask'
|
||||
# 0: background, 1: foreground
|
||||
mask = np.clip(input_mask, 0, 1)
|
||||
contour_radius = (contour_width - 1) // 2
|
||||
|
||||
dist_transform_fore = cv2.distanceTransform(mask, cv2.DIST_L2, 3)
|
||||
dist_transform_back = cv2.distanceTransform(1-mask, cv2.DIST_L2, 3)
|
||||
dist_map = dist_transform_fore - dist_transform_back
|
||||
# ...:::!!!:::...
|
||||
contour_radius += 2
|
||||
contour_mask = np.abs(np.clip(dist_map, -contour_radius, contour_radius))
|
||||
contour_mask = contour_mask / np.max(contour_mask)
|
||||
contour_mask[contour_mask>0.5] = 1.
|
||||
|
||||
# paint mask
|
||||
painted_image = vis_add_mask(input_image.copy(), mask.copy(), mask_color, mask_alpha)
|
||||
# paint contour
|
||||
painted_image = vis_add_mask(painted_image.copy(), 1-contour_mask, contour_color, 1)
|
||||
|
||||
return painted_image
|
||||
|
||||
def background_remover(input_image, input_mask):
|
||||
"""
|
||||
input_image: H, W, 3, np.array
|
||||
input_mask: H, W, np.array
|
||||
|
||||
image_wo_background: PIL.Image
|
||||
"""
|
||||
assert input_image.shape[:2] == input_mask.shape, 'different shape between image and mask'
|
||||
# 0: background, 1: foreground
|
||||
mask = np.expand_dims(np.clip(input_mask, 0, 1), axis=2)*255
|
||||
image_wo_background = np.concatenate([input_image, mask], axis=2) # H, W, 4
|
||||
image_wo_background = Image.fromarray(image_wo_background).convert('RGBA')
|
||||
|
||||
return image_wo_background
|
||||
|
||||
if __name__ == '__main__':
|
||||
input_image = np.array(Image.open('images/painter_input_image.jpg').convert('RGB'))
|
||||
input_mask = np.array(Image.open('images/painter_input_mask.jpg').convert('P'))
|
||||
|
||||
# example of mask painter
|
||||
mask_color = 3
|
||||
mask_alpha = 0.7
|
||||
contour_color = 1
|
||||
contour_width = 5
|
||||
|
||||
# save
|
||||
painted_image = Image.fromarray(input_image)
|
||||
painted_image.save('images/original.png')
|
||||
|
||||
painted_image = mask_painter(input_image, input_mask, mask_color, mask_alpha, contour_color, contour_width)
|
||||
# save
|
||||
painted_image = Image.fromarray(input_image)
|
||||
painted_image.save('images/original1.png')
|
||||
|
||||
# example of point painter
|
||||
input_image = np.array(Image.open('images/painter_input_image.jpg').convert('RGB'))
|
||||
input_points = np.array([[500, 375], [70, 600]]) # x, y
|
||||
point_color = 5
|
||||
point_alpha = 0.9
|
||||
point_radius = 15
|
||||
contour_color = 2
|
||||
contour_width = 5
|
||||
painted_image_1 = point_painter(input_image, input_points, point_color, point_alpha, point_radius, contour_color, contour_width)
|
||||
# save
|
||||
painted_image = Image.fromarray(painted_image_1)
|
||||
painted_image.save('images/point_painter_1.png')
|
||||
|
||||
input_image = np.array(Image.open('images/painter_input_image.jpg').convert('RGB'))
|
||||
painted_image_2 = point_painter(input_image, input_points, point_color=9, point_radius=20, contour_color=29)
|
||||
# save
|
||||
painted_image = Image.fromarray(painted_image_2)
|
||||
painted_image.save('images/point_painter_2.png')
|
||||
|
||||
# example of background remover
|
||||
input_image = np.array(Image.open('images/original.png').convert('RGB'))
|
||||
image_wo_background = background_remover(input_image, input_mask) # return PIL.Image
|
||||
image_wo_background.save('images/image_wo_background.png')
|
||||
@@ -0,0 +1,104 @@
|
||||
import os
|
||||
import argparse
|
||||
from safetensors.torch import load_file
|
||||
import torch
|
||||
from .objectclear.pipelines import ObjectClearPipeline
|
||||
from .objectclear.utils import resize_by_short_side
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
import gc
|
||||
|
||||
|
||||
def loader_objectclear(cf_model,cf_vae,cur_path,device,cache_dir,use_fp16,postfuse_path):
|
||||
# ------------------ set up ObjectClear pipeline -------------------
|
||||
torch_dtype = torch.float16 if use_fp16 else torch.float32
|
||||
variant = "fp16" if use_fp16 else None
|
||||
UNET = load_singlefile(cf_model,cur_path,device)
|
||||
VAE = load_singlevae(cf_vae,cur_path,device)
|
||||
|
||||
pipe = ObjectClearPipeline.from_pretrained_with_custom_modules(
|
||||
os.path.join(cur_path,"sdxl_repo"),
|
||||
unet=UNET,
|
||||
vae=VAE,
|
||||
torch_dtype=torch_dtype,
|
||||
apply_attention_guided_fusion=True,
|
||||
cache_dir=cache_dir,
|
||||
variant=variant,
|
||||
postfuse_path=postfuse_path,
|
||||
|
||||
)
|
||||
pipe.to(device)
|
||||
#pipe.enable_model_cpu_offload()
|
||||
return pipe
|
||||
|
||||
|
||||
|
||||
def inference_objectclear(pipe,input_img_list, input_mask_list,device,positive,negative,steps=20,seed=42,strength=0.99,guidance_scale=2.5,object_embeds=None,short_size=512):
|
||||
# -------------------- start to processing ---------------------
|
||||
output_img_list = []
|
||||
for i, (image, mask) in enumerate(zip(input_img_list, input_mask_list)):
|
||||
print(f'[{i+1}/{len(input_img_list)}] Processing')
|
||||
|
||||
image = resize_by_short_side(image, short_size, resample=Image.BICUBIC)
|
||||
mask = resize_by_short_side(mask, short_size, resample=Image.NEAREST)
|
||||
|
||||
|
||||
w, h = image.size
|
||||
result = pipe(
|
||||
prompt=None,#"remove the instance of object",
|
||||
image=image,
|
||||
mask_image=mask,
|
||||
generator=torch.Generator(device=device).manual_seed(seed),
|
||||
num_inference_steps=steps,
|
||||
strength=strength,
|
||||
guidance_scale=guidance_scale,
|
||||
height=h,
|
||||
width=w,
|
||||
return_attn_map=False,
|
||||
prompt_embeds= positive[0][0].to(device,dtype=torch.float16),
|
||||
negative_prompt_embeds= negative[0][0].to(device,dtype=torch.float16),
|
||||
pooled_prompt_embeds= positive[0][1]["pooled_output"].to(device, dtype=torch.float16),
|
||||
negative_pooled_prompt_embeds= negative[0][1]["pooled_output"].to(device, dtype=torch.float16),
|
||||
object_embeds=object_embeds[i].to(device, dtype=torch.float16),
|
||||
)
|
||||
|
||||
fused_img_pil = result.images[0]
|
||||
output_img_list.append(fused_img_pil)
|
||||
|
||||
# save results
|
||||
# save_path = os.path.join(result_root, f'{basename}.png')
|
||||
# fused_img_pil = fused_img_pil.resize(image_or.size)
|
||||
# fused_img_pil.save(save_path)
|
||||
|
||||
#print(f'\nAll results are saved in {result_root}')
|
||||
return output_img_list
|
||||
|
||||
|
||||
|
||||
def load_singlefile(model,cur_path,device):
|
||||
from diffusers import UNet2DConditionModel
|
||||
|
||||
config_file = os.path.join(cur_path,"sdxl_repo/unet/config.json")
|
||||
unet_state_dict=load_file(model)
|
||||
unet_config = UNet2DConditionModel.load_config(config_file)
|
||||
Unet = UNet2DConditionModel.from_config(unet_config).to(device, torch.float16)
|
||||
Unet.load_state_dict(unet_state_dict, strict=False)
|
||||
del unet_state_dict
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
return Unet
|
||||
|
||||
|
||||
def load_singlevae(VAE,cur_path,device):
|
||||
from diffusers import AutoencoderKL
|
||||
|
||||
vae_config = os.path.join(cur_path, "sdxl_repo/vae/config.json")
|
||||
vae_state_dict=load_file(VAE)
|
||||
ae_config = AutoencoderKL.load_config(vae_config)
|
||||
AE = AutoencoderKL.from_config(ae_config).to(device)
|
||||
AE.load_state_dict(vae_state_dict, strict=False)
|
||||
del vae_state_dict
|
||||
torch.cuda.empty_cache()
|
||||
return AE
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 584 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 334 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.6 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 1.6 MiB |
Binary file not shown.
|
After Width: | Height: | Size: 670 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 9.4 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 17 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 16 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 15 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 24 KiB |
+298
@@ -0,0 +1,298 @@
|
||||
# !/usr/bin/env python
|
||||
# -*- coding: UTF-8 -*-
|
||||
import os
|
||||
import torch
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
import cv2
|
||||
import gc
|
||||
|
||||
from comfy.utils import common_upscale,ProgressBar
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
cur_path = os.path.dirname(os.path.abspath(__file__))
|
||||
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
|
||||
|
||||
|
||||
def gc_cleanup():
|
||||
"""
|
||||
清理内存
|
||||
"""
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def find_gaussian_files(base_path):
|
||||
found_files = []
|
||||
for root, dirs, files in os.walk(base_path):
|
||||
for file in files:
|
||||
if file.endswith('.pth') and 'gaussian' in file:
|
||||
found_files.append(os.path.join(root, file))
|
||||
return found_files
|
||||
|
||||
def cv2pil(cv_image):
|
||||
"""
|
||||
将OpenCV图像转换为PIL图像
|
||||
:param cv_image: OpenCV图像
|
||||
:return: PIL图像
|
||||
"""
|
||||
# 将图像从BGR转换为RGB
|
||||
rgb_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB)
|
||||
# 使用PIL的Image.fromarray方法将NumPy数组转换为PIL图像
|
||||
pil_image = Image.fromarray(rgb_image)
|
||||
return pil_image
|
||||
|
||||
|
||||
|
||||
def tensor_to_pil(tensor):
|
||||
image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
|
||||
image = Image.fromarray(image_np, mode='RGB')
|
||||
return image
|
||||
|
||||
def tensor2pil_list(image):
|
||||
B,_,_,_=image.size()
|
||||
if B==1:
|
||||
ref_image_list=[tensor_to_pil(image)]
|
||||
else:
|
||||
img_list = list(torch.chunk(image, chunks=B))
|
||||
ref_image_list = [tensor_to_pil(img) for img in img_list]
|
||||
return ref_image_list
|
||||
|
||||
|
||||
def mask2pil_list(mask):
|
||||
B,_,_=mask.size()
|
||||
if B==1:
|
||||
ref_image_list = [mask2pil(mask)]
|
||||
else:
|
||||
img_list = list(torch.chunk(mask, chunks=B))
|
||||
ref_image_list = [mask2pil(img) for img in img_list]
|
||||
return ref_image_list
|
||||
|
||||
def tensor2pil_list_u(image,width,height):
|
||||
B,_,_,_=image.size()
|
||||
if B==1:
|
||||
ref_image_list=[tensor2pil_upscale(image,width,height)]
|
||||
else:
|
||||
img_list = list(torch.chunk(image, chunks=B))
|
||||
ref_image_list = [tensor2pil_upscale(img,width,height) for img in img_list]
|
||||
return ref_image_list
|
||||
|
||||
|
||||
def mask2pil_list_u(mask,width,height):
|
||||
B,_,_=mask.size()
|
||||
if B==1:
|
||||
ref_image_list = [mask2pil_u(mask,width,height)]
|
||||
else:
|
||||
img_list = list(torch.chunk(mask, chunks=B))
|
||||
ref_image_list = [mask2pil_u(img,width,height) for img in img_list]
|
||||
return ref_image_list
|
||||
|
||||
def mask2pil(mask):
|
||||
upscaled_mask =mask
|
||||
mask_np = upscaled_mask.squeeze(0).mul(255).clamp(0, 255).byte().numpy()
|
||||
return Image.fromarray(mask_np, mode='L')
|
||||
|
||||
def mask2pil_u(mask,width,height):
|
||||
upscaled_mask = tensor_upscale(mask.unsqueeze(-1), width, height)
|
||||
mask_np = upscaled_mask.squeeze(0).squeeze(-1).mul(255).clamp(0, 255).byte().numpy()
|
||||
return Image.fromarray(mask_np, mode='L')
|
||||
|
||||
def tensor2pil_RGBA_list(image,mask,width,height):
|
||||
B,_,_,_=image.size()
|
||||
# 处理mask维度
|
||||
if len(mask.shape) == 2:
|
||||
mask = mask.unsqueeze(0) # 从 [H,W] 变为 [1,H,W]
|
||||
if mask.shape[1] == 64:
|
||||
return tensor2pil_list(image,width,height) #[1,64,64]
|
||||
mask = 1.0 - mask
|
||||
if B == 1:
|
||||
upscaled_img = tensor_upscale(image, width, height)
|
||||
upscaled_mask = tensor_upscale(mask.unsqueeze(-1), width, height) # 添加通道维度
|
||||
|
||||
# 转换为RGBA
|
||||
rgb = upscaled_img.squeeze(0).mul(255).clamp(0, 255).byte().numpy()
|
||||
alpha = upscaled_mask.squeeze(0).mul(255).clamp(0, 255).byte().numpy()
|
||||
rgba = np.dstack((rgb, alpha)) # 合并为RGBA
|
||||
ref_image_list = [Image.fromarray(rgba, mode='RGBA')]
|
||||
else:
|
||||
img_list = list(torch.chunk(image, chunks=B))
|
||||
mask_list = list(torch.chunk(mask, chunks=B))
|
||||
ref_image_list = []
|
||||
for img, msk in zip(img_list, mask_list):
|
||||
# 上采样每个图像和对应的mask
|
||||
upscaled_img = tensor_upscale(img, width, height)
|
||||
upscaled_mask = tensor_upscale(msk.unsqueeze(-1), width, height)
|
||||
|
||||
# 转换为RGBA -
|
||||
rgb = upscaled_img.squeeze(0).mul(255).clamp(0, 255).byte().numpy()
|
||||
alpha = upscaled_mask.squeeze(0).mul(255).clamp(0, 255).byte().numpy()
|
||||
rgba = np.dstack((rgb, alpha))
|
||||
ref_image_list.append(Image.fromarray(rgba, mode='RGBA'))
|
||||
return ref_image_list
|
||||
|
||||
|
||||
|
||||
def tensor_upscale(img_tensor, width, height):
|
||||
samples = img_tensor.movedim(-1, 1)
|
||||
img = common_upscale(samples, width, height, "nearest-exact", "center")
|
||||
samples = img.movedim(1, -1)
|
||||
return samples
|
||||
|
||||
def tensor2pil_upscale(img_tensor, width, height):
|
||||
samples = img_tensor.movedim(-1, 1)
|
||||
img = common_upscale(samples, width, height, "nearest-exact", "center")
|
||||
samples = img.movedim(1, -1)
|
||||
img_pil = tensor_to_pil(samples)
|
||||
return img_pil
|
||||
|
||||
|
||||
def tensor2cv(tensor_image,RGB2BGR=True):
|
||||
if len(tensor_image.shape)==4:#bhwc to hwc
|
||||
tensor_image=tensor_image.squeeze(0)
|
||||
if tensor_image.is_cuda:
|
||||
tensor_image = tensor_image.cpu().detach()
|
||||
tensor_image=tensor_image.numpy()
|
||||
#反归一化
|
||||
maxValue=tensor_image.max()
|
||||
tensor_image=tensor_image*255/maxValue
|
||||
img_cv2=np.uint8(tensor_image)#32 to uint8
|
||||
if RGB2BGR:
|
||||
img_cv2=cv2.cvtColor(img_cv2,cv2.COLOR_RGB2BGR)
|
||||
return img_cv2
|
||||
|
||||
def cvargb2tensor(img):
|
||||
assert type(img) == np.ndarray, 'the img type is {}, but ndarry expected'.format(type(img))
|
||||
img = torch.from_numpy(img.transpose((2, 0, 1)))
|
||||
return img.float().div(255).unsqueeze(0) # 255也可以改为256
|
||||
|
||||
def cv2tensor(img):
|
||||
assert type(img) == np.ndarray, 'the img type is {}, but ndarry expected'.format(type(img))
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
img = torch.from_numpy(img.transpose((2, 0, 1)))
|
||||
return img.float().div(255).unsqueeze(0) # 255也可以改为256
|
||||
|
||||
def images_generator(img_list: list,):
|
||||
#get img size
|
||||
sizes = {}
|
||||
for image_ in img_list:
|
||||
if isinstance(image_,Image.Image):
|
||||
count = sizes.get(image_.size, 0)
|
||||
sizes[image_.size] = count + 1
|
||||
elif isinstance(image_,np.ndarray):
|
||||
count = sizes.get(image_.shape[:2][::-1], 0)
|
||||
sizes[image_.shape[:2][::-1]] = count + 1
|
||||
else:
|
||||
raise "unsupport image list,must be pil or cv2!!!"
|
||||
size = max(sizes.items(), key=lambda x: x[1])[0]
|
||||
yield size[0], size[1]
|
||||
|
||||
# any to tensor
|
||||
def load_image(img_in):
|
||||
if isinstance(img_in, Image.Image):
|
||||
img_in=img_in.convert("RGB")
|
||||
i = np.array(img_in, dtype=np.float32)
|
||||
i = torch.from_numpy(i).div_(255)
|
||||
if i.shape[0] != size[1] or i.shape[1] != size[0]:
|
||||
i = torch.from_numpy(i).movedim(-1, 0).unsqueeze(0)
|
||||
i = common_upscale(i, size[0], size[1], "lanczos", "center")
|
||||
i = i.squeeze(0).movedim(0, -1).numpy()
|
||||
return i
|
||||
elif isinstance(img_in,np.ndarray):
|
||||
i=cv2.cvtColor(img_in,cv2.COLOR_BGR2RGB).astype(np.float32)
|
||||
i = torch.from_numpy(i).div_(255)
|
||||
#print(i.shape)
|
||||
return i
|
||||
else:
|
||||
raise "unsupport image list,must be pil,cv2 or tensor!!!"
|
||||
|
||||
total_images = len(img_list)
|
||||
processed_images = 0
|
||||
pbar = ProgressBar(total_images)
|
||||
images = map(load_image, img_list)
|
||||
try:
|
||||
prev_image = next(images)
|
||||
while True:
|
||||
next_image = next(images)
|
||||
yield prev_image
|
||||
processed_images += 1
|
||||
pbar.update_absolute(processed_images, total_images)
|
||||
prev_image = next_image
|
||||
except StopIteration:
|
||||
pass
|
||||
if prev_image is not None:
|
||||
yield prev_image
|
||||
|
||||
def load_images(img_list: list,):
|
||||
gen = images_generator(img_list)
|
||||
(width, height) = next(gen)
|
||||
images = torch.from_numpy(np.fromiter(gen, np.dtype((np.float32, (height, width, 3)))))
|
||||
if len(images) == 0:
|
||||
raise FileNotFoundError(f"No images could be loaded .")
|
||||
return images
|
||||
|
||||
def tensor2pil(tensor):
|
||||
image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
|
||||
image = Image.fromarray(image_np, mode='RGB')
|
||||
return image
|
||||
|
||||
def pil2narry(img):
|
||||
narry = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0)
|
||||
return narry
|
||||
|
||||
def equalize_lists(list1, list2):
|
||||
"""
|
||||
比较两个列表的长度，如果不一致，则将较短的列表复制以匹配较长列表的长度。
|
||||
|
||||
参数:
|
||||
list1 (list): 第一个列表
|
||||
list2 (list): 第二个列表
|
||||
|
||||
返回:
|
||||
tuple: 包含两个长度相等的列表的元组
|
||||
"""
|
||||
len1 = len(list1)
|
||||
len2 = len(list2)
|
||||
|
||||
if len1 == len2:
|
||||
pass
|
||||
elif len1 < len2:
|
||||
print("list1 is shorter than list2, copying list1 to match list2's length.")
|
||||
list1.extend(list1 * ((len2 // len1) + 1)) # 复制list1以匹配list2的长度
|
||||
list1 = list1[:len2] # 确保长度一致
|
||||
else:
|
||||
print("list2 is shorter than list1, copying list2 to match list1's length.")
|
||||
list2.extend(list2 * ((len1 // len2) + 1)) # 复制list2以匹配list1的长度
|
||||
list2 = list2[:len1] # 确保长度一致
|
||||
|
||||
return list1, list2
|
||||
|
||||
def file_exists(directory, filename):
|
||||
# 构建文件的完整路径
|
||||
file_path = os.path.join(directory, filename)
|
||||
# 检查文件是否存在
|
||||
return os.path.isfile(file_path)
|
||||
|
||||
def download_weights(file_dir,repo_id,subfolder="",pt_name=""):
|
||||
if subfolder:
|
||||
file_path = os.path.join(file_dir,subfolder, pt_name)
|
||||
sub_dir=os.path.join(file_dir,subfolder)
|
||||
if not os.path.exists(sub_dir):
|
||||
os.makedirs(sub_dir)
|
||||
if not os.path.exists(file_path):
|
||||
file_path = hf_hub_download(
|
||||
repo_id=repo_id,
|
||||
subfolder=subfolder,
|
||||
filename=pt_name,
|
||||
local_dir = file_dir,
|
||||
)
|
||||
return file_path
|
||||
else:
|
||||
file_path = os.path.join(file_dir, pt_name)
|
||||
if not os.path.exists(file_dir):
|
||||
os.makedirs(file_dir)
|
||||
if not os.path.exists(file_path):
|
||||
file_path = hf_hub_download(
|
||||
repo_id=repo_id,
|
||||
filename=pt_name,
|
||||
local_dir=file_dir,
|
||||
)
|
||||
return file_path
|
||||
@@ -0,0 +1,12 @@
|
||||
from .models import CLIPImageEncoder, PostfuseModule
|
||||
from .pipelines import ObjectClearPipeline
|
||||
from .utils import attention_guided_fusion, resize_by_short_side
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CLIPImageEncoder",
|
||||
"PostfuseModule",
|
||||
"ObjectClearPipeline",
|
||||
"attention_guided_fusion",
|
||||
"resize_by_short_side",
|
||||
]
|
||||
Binary file not shown.
@@ -0,0 +1,5 @@
|
||||
from .clip_image_encoder import CLIPImageEncoder
|
||||
from .postfuse_module import PostfuseModule
|
||||
|
||||
|
||||
__all__ = ["CLIPImageEncoder", "PostfuseModule"]
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,58 @@
|
||||
import torch.nn.functional as F
|
||||
import torchvision.transforms as T
|
||||
from transformers.models.clip.modeling_clip import (
|
||||
CLIPPreTrainedModel,
|
||||
CLIPModel,
|
||||
)
|
||||
|
||||
|
||||
class CLIPImageEncoder(CLIPPreTrainedModel):
|
||||
@staticmethod
|
||||
def from_pretrained(
|
||||
global_model_name_or_path,
|
||||
cache_dir
|
||||
):
|
||||
model = CLIPModel.from_pretrained(
|
||||
global_model_name_or_path,
|
||||
subfolder="image_prompt_encoder",
|
||||
cache_dir=cache_dir
|
||||
)
|
||||
vision_model = model.vision_model
|
||||
visual_projection = model.visual_projection
|
||||
vision_processor = T.Normalize(
|
||||
(0.48145466, 0.4578275, 0.40821073),
|
||||
(0.26862954, 0.26130258, 0.27577711),
|
||||
)
|
||||
return CLIPImageEncoder(
|
||||
vision_model,
|
||||
visual_projection,
|
||||
vision_processor,
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vision_model,
|
||||
visual_projection,
|
||||
vision_processor,
|
||||
):
|
||||
super().__init__(vision_model.config)
|
||||
self.vision_model = vision_model
|
||||
self.visual_projection = visual_projection
|
||||
self.vision_processor = vision_processor
|
||||
|
||||
self.image_size = vision_model.config.image_size
|
||||
|
||||
def forward(self, object_pixel_values):
|
||||
b, c, h, w = object_pixel_values.shape
|
||||
|
||||
if h != self.image_size or w != self.image_size:
|
||||
h, w = self.image_size, self.image_size
|
||||
object_pixel_values = F.interpolate(
|
||||
object_pixel_values, (h, w), mode="bilinear", antialias=True
|
||||
)
|
||||
|
||||
object_pixel_values = self.vision_processor(object_pixel_values)
|
||||
object_embeds = self.vision_model(object_pixel_values)[1]
|
||||
object_embeds = self.visual_projection(object_embeds)
|
||||
object_embeds = object_embeds.view(b, 1, -1)
|
||||
return object_embeds
|
||||
@@ -0,0 +1,56 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class PostfuseModule(nn.Module):
|
||||
def __init__(self, embed_dim, embed_dim_img):
|
||||
super().__init__()
|
||||
self.mlp1 = MLP(embed_dim_img, embed_dim, embed_dim, use_residual=False)
|
||||
self.mlp2 = MLP(embed_dim, embed_dim, embed_dim, use_residual=True)
|
||||
self.layer_norm = nn.LayerNorm(embed_dim)
|
||||
|
||||
@property
|
||||
def dtype(self):
|
||||
try:
|
||||
return next(self.parameters()).dtype
|
||||
except StopIteration:
|
||||
return torch.float32
|
||||
|
||||
def fuse_fn(self, object_embeds):
|
||||
text_object_embeds = self.mlp1(object_embeds)
|
||||
text_object_embeds = self.mlp2(text_object_embeds)
|
||||
text_object_embeds = self.layer_norm(text_object_embeds)
|
||||
return text_object_embeds
|
||||
|
||||
def forward(
|
||||
self,
|
||||
text_embeds,
|
||||
object_embeds,
|
||||
fuse_index,
|
||||
) -> torch.Tensor:
|
||||
text_object_embed = self.fuse_fn(object_embeds)
|
||||
text_embeds_new = text_embeds.clone()
|
||||
text_embeds_new[:, fuse_index, :] = text_object_embed.squeeze(1)
|
||||
|
||||
return text_embeds_new
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, in_dim, out_dim, hidden_dim, use_residual=True):
|
||||
super().__init__()
|
||||
if use_residual:
|
||||
assert in_dim == out_dim
|
||||
self.layernorm = nn.LayerNorm(in_dim)
|
||||
self.fc1 = nn.Linear(in_dim, hidden_dim)
|
||||
self.fc2 = nn.Linear(hidden_dim, out_dim)
|
||||
self.use_residual = use_residual
|
||||
self.act_fn = nn.GELU()
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
x = self.layernorm(x)
|
||||
x = self.fc1(x)
|
||||
x = self.act_fn(x)
|
||||
x = self.fc2(x)
|
||||
if self.use_residual:
|
||||
x = x + residual
|
||||
return x
|
||||
@@ -0,0 +1,4 @@
|
||||
from .pipeline_objectclear import ObjectClearPipeline
|
||||
|
||||
|
||||
__all__ = ["ObjectClearPipeline"]
|
||||
Binary file not shown.
Binary file not shown.
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,5 @@
|
||||
from .attention_guided_fusion import attention_guided_fusion
|
||||
from .image_utils import pad_to_multiple, crop_to_original, resize_by_short_side
|
||||
|
||||
|
||||
__all__ = ["attention_guided_fusion", "pad_to_multiple", "crop_to_original", "resize_by_short_side"]
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,78 @@
|
||||
import numpy as np
|
||||
import cv2
|
||||
from scipy.ndimage import convolve, zoom
|
||||
from .image_utils import pad_to_multiple, crop_to_original
|
||||
|
||||
|
||||
def wavelet_blur_np(image: np.ndarray, radius: int):
|
||||
kernel = np.array([
|
||||
[0.0625, 0.125, 0.0625],
|
||||
[0.125, 0.25, 0.125],
|
||||
[0.0625, 0.125, 0.0625]
|
||||
], dtype=np.float32)
|
||||
|
||||
blurred = np.empty_like(image)
|
||||
for c in range(image.shape[0]):
|
||||
blurred_c = convolve(image[c], kernel, mode='nearest')
|
||||
if radius > 1:
|
||||
blurred_c = zoom(zoom(blurred_c, 1 / radius, order=1), radius, order=1)
|
||||
blurred[c] = blurred_c
|
||||
return blurred
|
||||
|
||||
def wavelet_decomposition_np(image: np.ndarray, levels=5):
|
||||
high_freq = np.zeros_like(image)
|
||||
for i in range(levels):
|
||||
radius = 2 ** i
|
||||
low_freq = wavelet_blur_np(image, radius)
|
||||
high_freq += (image - low_freq)
|
||||
image = low_freq
|
||||
return high_freq, low_freq
|
||||
|
||||
def wavelet_reconstruction_np(content_feat: np.ndarray, style_feat: np.ndarray):
|
||||
content_high, _ = wavelet_decomposition_np(content_feat)
|
||||
_, style_low = wavelet_decomposition_np(style_feat)
|
||||
return content_high + style_low
|
||||
|
||||
def wavelet_color_fix_np(fused: np.ndarray, mask: np.ndarray) -> np.ndarray:
|
||||
fused_np = fused.astype(np.float32) / 255.0
|
||||
mask_np = mask.astype(np.float32) / 255.0
|
||||
|
||||
fused_np = fused_np.transpose(2, 0, 1)
|
||||
mask_np = mask_np.transpose(2, 0, 1)
|
||||
|
||||
result_np = wavelet_reconstruction_np(fused_np, mask_np)
|
||||
|
||||
result_np = result_np.transpose(1, 2, 0)
|
||||
result_np = np.clip(result_np * 255.0, 0, 255).astype(np.uint8)
|
||||
|
||||
return result_np
|
||||
|
||||
def attention_guided_fusion(ori: np.ndarray, removed: np.ndarray, attn_map: np.ndarray, multiple: int = 8):
|
||||
H, W = ori.shape[:2]
|
||||
attn_map = attn_map.astype(np.float32)
|
||||
_, attn_map = cv2.threshold(attn_map, 128, 255, cv2.THRESH_BINARY)
|
||||
am = attn_map.astype(np.float32)
|
||||
am = am/255.0
|
||||
am_up = cv2.resize(am, (W, H), interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (21,21))
|
||||
am_d = cv2.dilate(am_up, kernel, iterations=1)
|
||||
am_d = cv2.GaussianBlur(am_d.astype(np.float32), (9,9), sigmaX=2)
|
||||
|
||||
am_merged = np.maximum(am_up, am_d)
|
||||
am_merged = np.clip(am_merged, 0, 1)
|
||||
|
||||
attn_up_3c = np.stack([am_merged]*3, axis=-1)
|
||||
attn_up_ori_3c = np.stack([am_up]*3, axis=-1)
|
||||
|
||||
ori_out = ori * (1 - attn_up_ori_3c)
|
||||
rem_out = removed * (1 - attn_up_ori_3c)
|
||||
|
||||
ori_pad, h0, w0 = pad_to_multiple(ori_out, multiple)
|
||||
rem_pad, _, _ = pad_to_multiple(rem_out, multiple)
|
||||
|
||||
wave_rgb = wavelet_color_fix_np(ori_pad, rem_pad)
|
||||
wave = crop_to_original(wave_rgb, h0, w0)
|
||||
|
||||
fused = (wave * (1 - attn_up_3c) + removed * attn_up_3c).astype(np.uint8)
|
||||
return fused
|
||||
@@ -0,0 +1,33 @@
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def pad_to_multiple(image: np.ndarray, multiple: int = 8):
|
||||
h, w = image.shape[:2]
|
||||
pad_h = (multiple - h % multiple) % multiple
|
||||
pad_w = (multiple - w % multiple) % multiple
|
||||
if image.ndim == 3:
|
||||
padded = np.pad(image, ((0, pad_h), (0, pad_w), (0,0)), mode='reflect')
|
||||
else:
|
||||
padded = np.pad(image, ((0, pad_h), (0, pad_w)), mode='reflect')
|
||||
return padded, h, w
|
||||
|
||||
def crop_to_original(image: np.ndarray, h: int, w: int):
|
||||
return image[:h, :w]
|
||||
|
||||
def resize_by_short_side(image, target_short=512, resample=Image.BICUBIC):
|
||||
w, h = image.size
|
||||
if min(w, h) < target_short:
|
||||
new_w = (w + 15) // 16 * 16
|
||||
new_h = (h + 15) // 16 * 16
|
||||
else:
|
||||
if w < h:
|
||||
new_w = target_short
|
||||
new_h = int(h * target_short / w)
|
||||
else:
|
||||
new_h = target_short
|
||||
new_w = int(w * target_short / h)
|
||||
new_w = (new_w + 15) // 16 * 16
|
||||
new_h = (new_h + 15) // 16 * 16
|
||||
|
||||
return image.resize((new_w, new_h), resample=resample)
|
||||
@@ -0,0 +1,16 @@
|
||||
[project]
|
||||
name = "objectclear"
|
||||
description = "ObjectClear:Complete Object Removal via Object-Effect Attention,you can try it in ComfyUI"
|
||||
version = "1.0.0"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = ["accelerate", "torch", "torchvision", "numpy", "opencv-python", "pillow", "transformers", "scipy", "diffusers"]
|
||||
|
||||
[project.urls]
|
||||
Repository = "https://github.com/smthemex/ComfyUI_ObjectClear"
|
||||
# Used by Comfy Registry https://registry.comfy.org
|
||||
|
||||
[tool.comfy]
|
||||
PublisherId = "smthemex"
|
||||
DisplayName = "ComfyUI_ObjectClear"
|
||||
Icon = ""
|
||||
includes = []
|
||||
@@ -0,0 +1,9 @@
|
||||
accelerate
|
||||
torch
|
||||
torchvision
|
||||
numpy
|
||||
opencv-python
|
||||
pillow
|
||||
transformers
|
||||
scipy
|
||||
diffusers
|
||||
@@ -0,0 +1,47 @@
|
||||
{
|
||||
"_name_or_path": "openai/clip-vit-large-patch14",
|
||||
"architectures": [
|
||||
"CLIPModel"
|
||||
],
|
||||
"initializer_factor": 1.0,
|
||||
"logit_scale_init_value": 2.6592,
|
||||
"model_type": "clip",
|
||||
"projection_dim": 768,
|
||||
"text_config": {
|
||||
"attention_dropout": 0.0,
|
||||
"dropout": 0.0,
|
||||
"hidden_act": "quick_gelu",
|
||||
"hidden_size": 768,
|
||||
"initializer_factor": 1.0,
|
||||
"initializer_range": 0.02,
|
||||
"intermediate_size": 3072,
|
||||
"layer_norm_eps": 1e-05,
|
||||
"max_position_embeddings": 77,
|
||||
"model_type": "clip_text_model",
|
||||
"num_attention_heads": 12,
|
||||
"num_hidden_layers": 12,
|
||||
"projection_dim": 768,
|
||||
"torch_dtype": "float32",
|
||||
"vocab_size": 49408
|
||||
},
|
||||
"torch_dtype": "float32",
|
||||
"transformers_version": "4.52.3",
|
||||
"vision_config": {
|
||||
"attention_dropout": 0.0,
|
||||
"dropout": 0.0,
|
||||
"hidden_act": "quick_gelu",
|
||||
"hidden_size": 1024,
|
||||
"image_size": 224,
|
||||
"initializer_factor": 1.0,
|
||||
"initializer_range": 0.02,
|
||||
"intermediate_size": 4096,
|
||||
"layer_norm_eps": 1e-05,
|
||||
"model_type": "clip_vision_model",
|
||||
"num_attention_heads": 16,
|
||||
"num_channels": 3,
|
||||
"num_hidden_layers": 24,
|
||||
"patch_size": 14,
|
||||
"projection_dim": 768,
|
||||
"torch_dtype": "float32"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
{
|
||||
"_class_name": "ObjectClearPipeline",
|
||||
"_diffusers_version": "0.21.0.dev0",
|
||||
"_name_or_path": "stabilityai/stable-diffusion-xl-base-1.0",
|
||||
"force_zeros_for_empty_prompt": true,
|
||||
"requires_aesthetics_score": false,
|
||||
"scheduler": [
|
||||
"diffusers",
|
||||
"EulerDiscreteScheduler"
|
||||
],
|
||||
"text_encoder": [
|
||||
"transformers",
|
||||
"CLIPTextModel"
|
||||
],
|
||||
"text_encoder_2": [
|
||||
"transformers",
|
||||
"CLIPTextModelWithProjection"
|
||||
],
|
||||
"tokenizer": [
|
||||
"transformers",
|
||||
"CLIPTokenizer"
|
||||
],
|
||||
"tokenizer_2": [
|
||||
"transformers",
|
||||
"CLIPTokenizer"
|
||||
],
|
||||
"unet": [
|
||||
"diffusers",
|
||||
"UNet2DConditionModel"
|
||||
],
|
||||
"vae": [
|
||||
"diffusers",
|
||||
"AutoencoderKL"
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"_class_name": "PostfuseModule",
|
||||
"embed_dim": 2048,
|
||||
"embed_dim_img": 768
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
{
|
||||
"_class_name": "EulerDiscreteScheduler",
|
||||
"_diffusers_version": "0.21.0.dev0",
|
||||
"beta_end": 0.012,
|
||||
"beta_schedule": "scaled_linear",
|
||||
"beta_start": 0.00085,
|
||||
"clip_sample": false,
|
||||
"interpolation_type": "linear",
|
||||
"num_train_timesteps": 1000,
|
||||
"prediction_type": "epsilon",
|
||||
"sample_max_value": 1.0,
|
||||
"set_alpha_to_one": false,
|
||||
"skip_prk_steps": true,
|
||||
"steps_offset": 1,
|
||||
"timestep_spacing": "leading",
|
||||
"trained_betas": null,
|
||||
"use_karras_sigmas": false
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
{
|
||||
"_name_or_path": "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
|
||||
"architectures": [
|
||||
"CLIPTextModel"
|
||||
],
|
||||
"attention_dropout": 0.0,
|
||||
"bos_token_id": 0,
|
||||
"dropout": 0.0,
|
||||
"eos_token_id": 2,
|
||||
"hidden_act": "quick_gelu",
|
||||
"hidden_size": 768,
|
||||
"initializer_factor": 1.0,
|
||||
"initializer_range": 0.02,
|
||||
"intermediate_size": 3072,
|
||||
"layer_norm_eps": 1e-05,
|
||||
"max_position_embeddings": 77,
|
||||
"model_type": "clip_text_model",
|
||||
"num_attention_heads": 12,
|
||||
"num_hidden_layers": 12,
|
||||
"pad_token_id": 1,
|
||||
"projection_dim": 768,
|
||||
"torch_dtype": "float16",
|
||||
"transformers_version": "4.29.2",
|
||||
"vocab_size": 49408
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
{
|
||||
"_name_or_path": "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
|
||||
"architectures": [
|
||||
"CLIPTextModelWithProjection"
|
||||
],
|
||||
"attention_dropout": 0.0,
|
||||
"bos_token_id": 0,
|
||||
"dropout": 0.0,
|
||||
"eos_token_id": 2,
|
||||
"hidden_act": "gelu",
|
||||
"hidden_size": 1280,
|
||||
"initializer_factor": 1.0,
|
||||
"initializer_range": 0.02,
|
||||
"intermediate_size": 5120,
|
||||
"layer_norm_eps": 1e-05,
|
||||
"max_position_embeddings": 77,
|
||||
"model_type": "clip_text_model",
|
||||
"num_attention_heads": 20,
|
||||
"num_hidden_layers": 32,
|
||||
"pad_token_id": 1,
|
||||
"projection_dim": 1280,
|
||||
"torch_dtype": "float16",
|
||||
"transformers_version": "4.29.2",
|
||||
"vocab_size": 49408
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,24 @@
|
||||
{
|
||||
"bos_token": {
|
||||
"content": "<|startoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"eos_token": {
|
||||
"content": "<|endoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"pad_token": "<|endoftext|>",
|
||||
"unk_token": {
|
||||
"content": "<|endoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
{
|
||||
"add_prefix_space": false,
|
||||
"bos_token": {
|
||||
"__type": "AddedToken",
|
||||
"content": "<|startoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"clean_up_tokenization_spaces": true,
|
||||
"do_lower_case": true,
|
||||
"eos_token": {
|
||||
"__type": "AddedToken",
|
||||
"content": "<|endoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"errors": "replace",
|
||||
"model_max_length": 77,
|
||||
"pad_token": "<|endoftext|>",
|
||||
"tokenizer_class": "CLIPTokenizer",
|
||||
"unk_token": {
|
||||
"__type": "AddedToken",
|
||||
"content": "<|endoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,24 @@
|
||||
{
|
||||
"bos_token": {
|
||||
"content": "<|startoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"eos_token": {
|
||||
"content": "<|endoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"pad_token": "!",
|
||||
"unk_token": {
|
||||
"content": "<|endoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
{
|
||||
"add_prefix_space": false,
|
||||
"bos_token": {
|
||||
"__type": "AddedToken",
|
||||
"content": "<|startoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"clean_up_tokenization_spaces": true,
|
||||
"do_lower_case": true,
|
||||
"eos_token": {
|
||||
"__type": "AddedToken",
|
||||
"content": "<|endoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"errors": "replace",
|
||||
"model_max_length": 77,
|
||||
"pad_token": "!",
|
||||
"tokenizer_class": "CLIPTokenizer",
|
||||
"unk_token": {
|
||||
"__type": "AddedToken",
|
||||
"content": "<|endoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,78 @@
|
||||
{
|
||||
"_class_name": "UNet2DConditionModel",
|
||||
"_diffusers_version": "0.21.0.dev0",
|
||||
"_name_or_path": "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
|
||||
"act_fn": "silu",
|
||||
"addition_embed_type": "text_time",
|
||||
"addition_embed_type_num_heads": 64,
|
||||
"addition_time_embed_dim": 256,
|
||||
"attention_head_dim": [
|
||||
5,
|
||||
10,
|
||||
20
|
||||
],
|
||||
"attention_type": "default",
|
||||
"block_out_channels": [
|
||||
320,
|
||||
640,
|
||||
1280
|
||||
],
|
||||
"center_input_sample": false,
|
||||
"class_embed_type": null,
|
||||
"class_embeddings_concat": false,
|
||||
"conv_in_kernel": 3,
|
||||
"conv_out_kernel": 3,
|
||||
"cross_attention_dim": 2048,
|
||||
"cross_attention_norm": null,
|
||||
"decay": 0.9999,
|
||||
"down_block_types": [
|
||||
"DownBlock2D",
|
||||
"CrossAttnDownBlock2D",
|
||||
"CrossAttnDownBlock2D"
|
||||
],
|
||||
"downsample_padding": 1,
|
||||
"dual_cross_attention": false,
|
||||
"encoder_hid_dim": null,
|
||||
"encoder_hid_dim_type": null,
|
||||
"flip_sin_to_cos": true,
|
||||
"freq_shift": 0,
|
||||
"in_channels": 9,
|
||||
"inv_gamma": 1.0,
|
||||
"layers_per_block": 2,
|
||||
"mid_block_only_cross_attention": null,
|
||||
"mid_block_scale_factor": 1,
|
||||
"mid_block_type": "UNetMidBlock2DCrossAttn",
|
||||
"min_decay": 0.0,
|
||||
"norm_eps": 1e-05,
|
||||
"norm_num_groups": 32,
|
||||
"num_attention_heads": null,
|
||||
"num_class_embeds": null,
|
||||
"only_cross_attention": false,
|
||||
"optimization_step": 37000,
|
||||
"out_channels": 4,
|
||||
"power": 0.6666666666666666,
|
||||
"projection_class_embeddings_input_dim": 2816,
|
||||
"resnet_out_scale_factor": 1.0,
|
||||
"resnet_skip_time_act": false,
|
||||
"resnet_time_scale_shift": "default",
|
||||
"sample_size": 128,
|
||||
"time_cond_proj_dim": null,
|
||||
"time_embedding_act_fn": null,
|
||||
"time_embedding_dim": null,
|
||||
"time_embedding_type": "positional",
|
||||
"timestep_post_act": null,
|
||||
"transformer_layers_per_block": [
|
||||
1,
|
||||
2,
|
||||
10
|
||||
],
|
||||
"up_block_types": [
|
||||
"CrossAttnUpBlock2D",
|
||||
"CrossAttnUpBlock2D",
|
||||
"UpBlock2D"
|
||||
],
|
||||
"upcast_attention": null,
|
||||
"update_after_step": 0,
|
||||
"use_ema_warmup": false,
|
||||
"use_linear_projection": true
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
{
|
||||
"_class_name": "AutoencoderKL",
|
||||
"_diffusers_version": "0.21.0.dev0",
|
||||
"_name_or_path": "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
|
||||
"act_fn": "silu",
|
||||
"block_out_channels": [
|
||||
128,
|
||||
256,
|
||||
512,
|
||||
512
|
||||
],
|
||||
"down_block_types": [
|
||||
"DownEncoderBlock2D",
|
||||
"DownEncoderBlock2D",
|
||||
"DownEncoderBlock2D",
|
||||
"DownEncoderBlock2D"
|
||||
],
|
||||
"force_upcast": false,
|
||||
"in_channels": 3,
|
||||
"latent_channels": 4,
|
||||
"layers_per_block": 2,
|
||||
"norm_num_groups": 32,
|
||||
"out_channels": 3,
|
||||
"sample_size": 512,
|
||||
"scaling_factor": 0.13025,
|
||||
"up_block_types": [
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D"
|
||||
]
|
||||
}
|
||||
Reference in New Issue
Block a user