75 lines
2.4 KiB
Python
75 lines
2.4 KiB
Python
import json
|
|
from pathlib import Path
|
|
|
|
import comfy.model_management as model_management
|
|
import torch
|
|
from comfy.clip_vision import Output, clip_preprocess
|
|
from PIL import Image
|
|
from torchvision.transforms.functional import pil_to_tensor
|
|
from tqdm import tqdm
|
|
|
|
|
|
def generate_clip_features_json(clip_vision, path_to_images_folder: Path,
|
|
output_json_path: Path):
|
|
clip_features = []
|
|
|
|
image_path_list = list(path_to_images_folder.glob('**/*.*'))
|
|
|
|
model_management.load_model_gpu(clip_vision.patcher)
|
|
print("--------- Generating embeddings for images ---------")
|
|
for image_path in tqdm(image_path_list):
|
|
image = Image.open(image_path)
|
|
image_embeds = get_image_clip_embeddings(clip_vision, image)
|
|
|
|
clip_features.append((str(image_path.name), image_embeds))
|
|
|
|
with open(output_json_path, 'w') as f:
|
|
json.dump(clip_features, f)
|
|
|
|
|
|
def get_image_clip_embeddings(clip_vision, image: Image):
|
|
model_management.load_model_gpu(clip_vision.patcher)
|
|
|
|
if type(image) is not torch.Tensor:
|
|
image = image_to_tensor(image)
|
|
image = image.to(clip_vision.load_device)
|
|
|
|
pixel_values = clip_preprocess(image)
|
|
|
|
out = clip_vision.model(pixel_values=pixel_values, intermediate_output=-2)
|
|
|
|
outputs = Output()
|
|
outputs["last_hidden_state"] = out[0].to(
|
|
model_management.intermediate_device())
|
|
outputs["image_embeds"] = out[2].to(
|
|
model_management.intermediate_device())
|
|
outputs["penultimate_hidden_states"] = out[1].to(
|
|
model_management.intermediate_device())
|
|
|
|
return outputs["image_embeds"].numpy().flatten().tolist()
|
|
|
|
|
|
def get_image_and_mask(file_name: str, path_to_images_folder: Path,
|
|
path_to_masks_folder: Path):
|
|
image = Image.open(path_to_images_folder / file_name)
|
|
mask_path = path_to_masks_folder / file_name
|
|
try:
|
|
mask = Image.open(mask_path.with_suffix('.png'))
|
|
except FileNotFoundError:
|
|
try:
|
|
mask = Image.open(mask_path.with_suffix('.jpg'))
|
|
except FileNotFoundError:
|
|
mask = Image.new('RGB', image.size, (0, 0, 0))
|
|
mask = mask.convert('RGB')
|
|
|
|
image = image_to_tensor(image)
|
|
mask = image_to_tensor(mask)
|
|
return image, mask
|
|
|
|
|
|
def image_to_tensor(image):
|
|
tensor = torch.clamp(pil_to_tensor(image).float() / 255., 0, 1)
|
|
tensor = tensor.unsqueeze(0)
|
|
tensor = tensor.permute(0, 2, 3, 1)
|
|
return tensor
|