456 lines
16 KiB
Python
456 lines
16 KiB
Python
import os
|
|
import sys
|
|
import json
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from PIL import Image
|
|
import numpy as np
|
|
from pathlib import Path
|
|
from typing import Dict, Optional
|
|
|
|
sys.path.append(os.path.dirname(__file__))
|
|
|
|
import folder_paths
|
|
|
|
from timm.data import resolve_data_config
|
|
from timm.data.transforms_factory import create_transform
|
|
from timm.models import create_model
|
|
|
|
from lsnet_model import lsnet_artist # noqa: F401
|
|
|
|
from inference_artist import (
|
|
load_checkpoint_state,
|
|
normalize_state_dict_keys,
|
|
resolve_num_classes,
|
|
resolve_feature_dim,
|
|
load_class_mapping
|
|
)
|
|
|
|
from sklearn.cluster import KMeans, DBSCAN, AgglomerativeClustering
|
|
from sklearn.manifold import TSNE
|
|
from sklearn.decomposition import PCA
|
|
import matplotlib.pyplot as plt
|
|
|
|
class LSNetModelLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
base_dir = os.path.join(folder_paths.models_dir, 'lsnet')
|
|
subfolders = []
|
|
if os.path.exists(base_dir):
|
|
subfolders = [f for f in os.listdir(base_dir) if os.path.isdir(os.path.join(base_dir, f))]
|
|
|
|
return {
|
|
"required": {
|
|
"model_folder": (subfolders, {"default": subfolders[0] if subfolders else ""}),
|
|
"device": ("STRING", {"default": "cuda"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LSNET_MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "LSNet"
|
|
|
|
def load(self, model_folder, device):
|
|
base_dir = os.path.join(folder_paths.models_dir, 'lsnet')
|
|
model_dir = os.path.join(base_dir, model_folder)
|
|
checkpoint_path = os.path.join(model_dir, "best_checkpoint.pth")
|
|
csv_path = os.path.join(model_dir, "class_mapping.csv")
|
|
|
|
if not os.path.exists(checkpoint_path):
|
|
raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
|
|
if not os.path.exists(csv_path):
|
|
raise FileNotFoundError(f"Class mapping CSV not found: {csv_path}")
|
|
class_mapping = load_class_mapping(csv_path)
|
|
state_dict = load_checkpoint_state(checkpoint_path)
|
|
state_dict = normalize_state_dict_keys(state_dict)
|
|
num_classes = resolve_num_classes(None, class_mapping, state_dict)
|
|
feature_dim = resolve_feature_dim(None, state_dict)
|
|
model = create_model(
|
|
'lsnet_xl_artist',
|
|
pretrained=False,
|
|
num_classes=num_classes,
|
|
feature_dim=feature_dim,
|
|
)
|
|
model.load_state_dict(state_dict, strict=False)
|
|
model.to(device)
|
|
model.eval()
|
|
config = resolve_data_config({}, model=model)
|
|
transform = create_transform(**config)
|
|
model_bundle = {
|
|
'model': model,
|
|
'transform': transform,
|
|
'class_mapping': class_mapping,
|
|
'device': device
|
|
}
|
|
|
|
return (model_bundle,)
|
|
|
|
class LSNetArtistInferenceNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"model": ("LSNET_MODEL",),
|
|
"top_k": ("INT", {"default": 5, "min": 1, "max": 100}),
|
|
"threshold": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("tag_string", "json_output")
|
|
FUNCTION = "process"
|
|
CATEGORY = "LSNet"
|
|
|
|
def process(self, image, model, top_k, threshold):
|
|
model_bundle = model
|
|
model = model_bundle['model']
|
|
transform = model_bundle['transform']
|
|
class_mapping = model_bundle['class_mapping']
|
|
device = model_bundle['device']
|
|
|
|
if image.ndim == 4:
|
|
image = image[0]
|
|
image = (image * 255).clamp(0, 255).byte().cpu().numpy()
|
|
pil_image = Image.fromarray(image)
|
|
|
|
# Preprocess image
|
|
image_tensor = transform(pil_image).unsqueeze(0) # Add batch dimension
|
|
|
|
# Classify
|
|
with torch.no_grad():
|
|
image_tensor = image_tensor.to(device)
|
|
logits = model(image_tensor, return_features=False)
|
|
probs = F.softmax(logits, dim=-1)
|
|
top_probs, top_indices = torch.topk(probs, k=min(top_k, probs.size(-1)), dim=-1)
|
|
|
|
results = []
|
|
for prob, idx in zip(top_probs[0].cpu().numpy(), top_indices[0].cpu().numpy()):
|
|
if prob >= threshold:
|
|
class_id = int(idx)
|
|
class_name = class_mapping.get(class_id, f"Class {class_id}")
|
|
results.append({
|
|
'class_id': class_id,
|
|
'class_name': class_name,
|
|
'probability': float(prob)
|
|
})
|
|
|
|
# Limit to top_k if more results after filtering
|
|
if len(results) > top_k:
|
|
results = results[:top_k]
|
|
|
|
# Prepare outputs
|
|
tags = [res['class_name'] for res in results]
|
|
tag_string = ",".join(tags)
|
|
tag_dict = {res['class_name']: res['probability'] for res in results}
|
|
json_output = json.dumps(tag_dict, ensure_ascii=False)
|
|
|
|
return (tag_string, json_output)
|
|
|
|
class LSNetArtistSimilarityNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"processed_image": ("IMAGE",),
|
|
"reference_images": ("IMAGE",),
|
|
"model": ("LSNET_MODEL",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("similarity_json",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "LSNet"
|
|
|
|
def process(self, processed_image, reference_images, model):
|
|
model_bundle = model
|
|
model = model_bundle['model']
|
|
transform = model_bundle['transform']
|
|
device = model_bundle['device']
|
|
|
|
def image_to_tensor(img):
|
|
if img.ndim == 4:
|
|
img = img[0]
|
|
img = (img * 255).clamp(0, 255).byte().cpu().numpy()
|
|
pil_img = Image.fromarray(img)
|
|
return transform(pil_img).unsqueeze(0)
|
|
|
|
processed_tensor = image_to_tensor(processed_image)
|
|
with torch.no_grad():
|
|
processed_tensor = processed_tensor.to(device)
|
|
processed_features = model(processed_tensor, return_features=True).cpu().numpy()[0]
|
|
|
|
references = []
|
|
similarities = []
|
|
num_refs = reference_images.shape[0] if reference_images.ndim == 4 else 1
|
|
for i in range(num_refs):
|
|
ref_img = reference_images[i] if reference_images.ndim == 4 else reference_images
|
|
ref_tensor = image_to_tensor(ref_img)
|
|
with torch.no_grad():
|
|
ref_tensor = ref_tensor.to(device)
|
|
ref_features = model(ref_tensor, return_features=True).cpu().numpy()[0]
|
|
references.append(ref_features.tolist())
|
|
sim = np.dot(processed_features, ref_features) / (np.linalg.norm(processed_features) * np.linalg.norm(ref_features))
|
|
similarities.append(float(sim))
|
|
|
|
result = {
|
|
"processed_features": processed_features.tolist(),
|
|
"reference_features": references,
|
|
"similarities": similarities
|
|
}
|
|
json_output = json.dumps(result, ensure_ascii=False)
|
|
|
|
return (json_output,)
|
|
|
|
class LSNetCommonFeaturesNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"reference_images": ("IMAGE",),
|
|
"model": ("LSNET_MODEL",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("TENSOR",)
|
|
RETURN_NAMES = ("common_features",)
|
|
FUNCTION = "process"
|
|
CATEGORY = "LSNet"
|
|
|
|
def process(self, reference_images, model):
|
|
model_bundle = model
|
|
model = model_bundle['model']
|
|
transform = model_bundle['transform']
|
|
device = model_bundle['device']
|
|
|
|
def image_to_tensor(img):
|
|
if img.ndim == 4:
|
|
img = img[0]
|
|
img = (img * 255).clamp(0, 255).byte().cpu().numpy()
|
|
pil_img = Image.fromarray(img)
|
|
return transform(pil_img).unsqueeze(0)
|
|
|
|
references = []
|
|
num_refs = reference_images.shape[0] if reference_images.ndim == 4 else 1
|
|
for i in range(num_refs):
|
|
ref_img = reference_images[i] if reference_images.ndim == 4 else reference_images
|
|
ref_tensor = image_to_tensor(ref_img)
|
|
with torch.no_grad():
|
|
ref_tensor = ref_tensor.to(device)
|
|
ref_features = model(ref_tensor, return_features=True).cpu().numpy()[0]
|
|
references.append(ref_features)
|
|
|
|
if references:
|
|
common_features = np.mean(np.array(references), axis=0)
|
|
else:
|
|
common_features = np.zeros(384)
|
|
return (torch.tensor(common_features),)
|
|
|
|
class LSNetClusteringNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"method": (["kmeans", "dbscan", "hierarchical"], {"default": "kmeans"}),
|
|
"n_clusters": ("INT", {"default": 10, "min": 2, "max": 100}),
|
|
"eps": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 10.0}),
|
|
"min_samples": ("INT", {"default": 5, "min": 1, "max": 50}),
|
|
"visualize": ("BOOLEAN", {"default": True}),
|
|
"viz_method": (["tsne", "pca"], {"default": "tsne"}),
|
|
"perplexity": ("INT", {"default": 30, "min": 5, "max": 100}),
|
|
},
|
|
"optional": {
|
|
"group_1": ("TENSOR",),
|
|
"group_2": ("TENSOR",),
|
|
"group_3": ("TENSOR",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "IMAGE")
|
|
RETURN_NAMES = ("clustering_json", "visualization")
|
|
FUNCTION = "cluster"
|
|
CATEGORY = "LSNet"
|
|
|
|
def cluster(self, method, n_clusters, eps, min_samples, visualize, viz_method, perplexity, group_1=None, group_2=None, group_3=None):
|
|
groups = []
|
|
group_sizes = []
|
|
for g in [group_1, group_2, group_3]:
|
|
if g is not None:
|
|
groups.append(g.cpu().numpy())
|
|
group_sizes.append(g.shape[0])
|
|
|
|
if not groups:
|
|
return (json.dumps({"error": "No groups provided"}), torch.zeros(1, 64, 64, 3))
|
|
|
|
features_np = np.vstack(groups)
|
|
if method == "kmeans":
|
|
clusterer = KMeans(n_clusters=n_clusters, random_state=42)
|
|
labels = clusterer.fit_predict(features_np)
|
|
centers = clusterer.cluster_centers_
|
|
elif method == "dbscan":
|
|
clusterer = DBSCAN(eps=eps, min_samples=min_samples)
|
|
labels = clusterer.fit_predict(features_np)
|
|
centers = None
|
|
elif method == "hierarchical":
|
|
clusterer = AgglomerativeClustering(n_clusters=n_clusters)
|
|
labels = clusterer.fit_predict(features_np)
|
|
centers = None
|
|
|
|
result = {
|
|
"method": method,
|
|
"n_samples": len(features_np),
|
|
"group_sizes": group_sizes,
|
|
"labels": labels.tolist(),
|
|
}
|
|
if centers is not None:
|
|
result["centers"] = centers.tolist()
|
|
|
|
json_output = json.dumps(result, ensure_ascii=False)
|
|
|
|
if visualize and len(features_np) > 1:
|
|
if viz_method == "tsne":
|
|
reducer = TSNE(n_components=2, perplexity=min(perplexity, len(features_np)-1), random_state=42)
|
|
else:
|
|
reducer = PCA(n_components=2, random_state=42)
|
|
|
|
reduced_features = reducer.fit_transform(features_np)
|
|
|
|
plt.figure(figsize=(10, 8))
|
|
unique_labels = np.unique(labels)
|
|
colors = plt.cm.rainbow(np.linspace(0, 1, len(unique_labels)))
|
|
|
|
for label, color in zip(unique_labels, colors):
|
|
mask = labels == label
|
|
plt.scatter(reduced_features[mask, 0], reduced_features[mask, 1],
|
|
color=color, label=f'Cluster {label}', alpha=0.7)
|
|
|
|
plt.title(f'{method.upper()} Clustering ({viz_method.upper()})')
|
|
plt.legend()
|
|
plt.tight_layout()
|
|
|
|
fig = plt.gcf()
|
|
fig.canvas.draw()
|
|
img_array = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
|
|
img_array = img_array.reshape(fig.canvas.get_width_height()[::-1] + (3,))
|
|
pil_image = Image.fromarray(img_array)
|
|
plt.close()
|
|
|
|
viz_tensor = torch.from_numpy(np.array(pil_image)).float() / 255.0
|
|
if viz_tensor.ndim == 3:
|
|
viz_tensor = viz_tensor.unsqueeze(0)
|
|
else:
|
|
viz_tensor = torch.zeros(1, 64, 64, 3)
|
|
|
|
return (json_output, viz_tensor)
|
|
|
|
class LSNetFeatureComparisonNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"model": ("LSNET_MODEL",),
|
|
},
|
|
"optional": {
|
|
"group_1": ("TENSOR",),
|
|
"group_2": ("TENSOR",),
|
|
"group_3": ("TENSOR",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("comparison_json",)
|
|
FUNCTION = "compare"
|
|
CATEGORY = "LSNet"
|
|
|
|
def compare(self, image, model, group_1=None, group_2=None, group_3=None):
|
|
model_bundle = model
|
|
model = model_bundle['model']
|
|
transform = model_bundle['transform']
|
|
device = model_bundle['device']
|
|
|
|
if image.ndim == 4:
|
|
image = image[0]
|
|
image_np = (image * 255).clamp(0, 255).byte().cpu().numpy()
|
|
pil_image = Image.fromarray(image_np)
|
|
image_tensor = transform(pil_image).unsqueeze(0)
|
|
|
|
with torch.no_grad():
|
|
image_tensor = image_tensor.to(device)
|
|
query_features = model(image_tensor, return_features=True).cpu().numpy()[0]
|
|
|
|
groups = []
|
|
for g in [group_1, group_2, group_3]:
|
|
if g is not None:
|
|
groups.append(g.cpu().numpy())
|
|
|
|
if not groups:
|
|
return (json.dumps({"error": "No groups provided"}),)
|
|
|
|
similarities = []
|
|
for group_feat in groups:
|
|
sim = np.dot(query_features, group_feat) / (np.linalg.norm(query_features) * np.linalg.norm(group_feat))
|
|
similarities.append(float(sim))
|
|
|
|
best_index = np.argmax(similarities)
|
|
best_similarity = similarities[best_index]
|
|
result = {
|
|
"best_group_index": int(best_index),
|
|
"best_similarity": best_similarity,
|
|
"all_similarities": similarities
|
|
}
|
|
json_output = json.dumps(result, ensure_ascii=False)
|
|
|
|
return (json_output,)
|
|
|
|
class LSNetArtistImageConnector:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image_1": ("IMAGE",),
|
|
"image_2": ("IMAGE",),
|
|
"image_3": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("stacked_images",)
|
|
FUNCTION = "connect"
|
|
CATEGORY = "LSNet"
|
|
|
|
def connect(self, image_1, image_2, image_3):
|
|
def normalize_image(img):
|
|
if img.ndim == 4:
|
|
img = img[0]
|
|
return img.unsqueeze(0)
|
|
|
|
img1 = normalize_image(image_1)
|
|
img2 = normalize_image(image_2)
|
|
img3 = normalize_image(image_3)
|
|
|
|
stacked = torch.cat([img1, img2, img3], dim=0)
|
|
return (stacked,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LSNetModelLoader": LSNetModelLoader,
|
|
"LSNetArtistInference": LSNetArtistInferenceNode,
|
|
"LSNetArtistSimilarity": LSNetArtistSimilarityNode,
|
|
"LSNetCommonFeatures": LSNetCommonFeaturesNode,
|
|
"LSNetClustering": LSNetClusteringNode,
|
|
"LSNetFeatureComparison": LSNetFeatureComparisonNode,
|
|
"LSNetArtistImageConnector": LSNetArtistImageConnector
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LSNetModelLoader": "LSNet Model Loader",
|
|
"LSNetArtistInference": "LSNet Artist Inference",
|
|
"LSNetArtistSimilarity": "LSNet Artist Similarity",
|
|
"LSNetCommonFeatures": "LSNet Common Features",
|
|
"LSNetClustering": "LSNet Clustering",
|
|
"LSNetFeatureComparison": "LSNet Feature Comparison",
|
|
"LSNetArtistImageConnector": "LSNet Image Connector"
|
|
}
|