155 lines
9.0 KiB
Python
155 lines
9.0 KiB
Python
"""The same analysis tab is mounted by WebUI and the independent Gradio app."""
|
|
import json
|
|
from functools import wraps
|
|
from pathlib import Path
|
|
import tempfile
|
|
|
|
import gradio as gr
|
|
from PIL import Image
|
|
|
|
from backend_lsnet.analysis import (extract_batch, thumbnail_batch, cache_bytes, read_cache,
|
|
analyze_cached, save_analysis, feature_tools, output_layout)
|
|
from backend_lsnet.model_paths import get_available_models, get_checkpoint_path
|
|
from feature_analysis import CHART_TYPES, TENSOR_LAYOUTS
|
|
from model_loading import FEATURE_OUTPUTS, load_model_bundle
|
|
|
|
|
|
def ui_errors(callback):
|
|
@wraps(callback)
|
|
def wrapped(*args, **kwargs):
|
|
try:
|
|
return callback(*args, **kwargs)
|
|
except (ValueError, FileNotFoundError, TypeError, KeyError) as error:
|
|
raise gr.Error(str(error)) from error
|
|
return wrapped
|
|
|
|
|
|
def file_path(value):
|
|
return str(value) if isinstance(value, (str, Path)) else value.name
|
|
|
|
|
|
@ui_errors
|
|
def extract_uploaded(files, model_name, device, output_type, layers, intermediate_norm, batch_size):
|
|
if not files:
|
|
raise gr.Error('请先上传图片。')
|
|
if len(files) > 512:
|
|
raise gr.Error('最多支持 512 张图片。')
|
|
images, labels = [], []
|
|
for item in files:
|
|
with Image.open(file_path(item)) as image:
|
|
images.append(image.copy())
|
|
labels.append(Path(file_path(item)).name)
|
|
bundle = load_model_bundle(checkpoint=get_checkpoint_path(model_name), device=device)
|
|
features = extract_batch(images, bundle, output_type, layers, intermediate_norm, int(batch_size))
|
|
cached = {'features': features, 'labels': labels, 'output_type': output_type, 'layers': layers,
|
|
'images': thumbnail_batch(images)}
|
|
directory = Path(tempfile.mkdtemp(prefix='kaloscope-features-'))
|
|
path = directory / 'features.npz'
|
|
path.write_bytes(cache_bytes(features, labels, output_type, layers))
|
|
return cached, f'已缓存 {len(images)} 张图片，TENSOR {list(features.shape)}。可反复制图，无需再推理。', str(path)
|
|
|
|
|
|
@ui_errors
|
|
def import_uploaded_cache(file):
|
|
if not file:
|
|
raise gr.Error('请选择 features.npz。')
|
|
cached = read_cache(file_path(file))
|
|
return cached, f'已导入 {len(cached["labels"])} 张图片的缓存，TENSOR {list(cached["features"].shape)}。未加载模型。'
|
|
|
|
|
|
@ui_errors
|
|
def plot_uploaded_cache(cached, chart_type, *values):
|
|
if cached is None:
|
|
raise gr.Error('先提取特征或导入缓存。')
|
|
options = dict(zip(ANALYSIS_OPTION_NAMES, values))
|
|
image, report, distances = analyze_cached(cached, chart_type, images=cached.get('images'), **options)
|
|
directory = tempfile.mkdtemp(prefix='kaloscope-analysis-')
|
|
files = save_analysis(directory, chart_type, image, report, distances)
|
|
return image[0].numpy(), report, files
|
|
|
|
|
|
@ui_errors
|
|
def tools_uploaded_cache(cached, operation, reference_index, groups, tensor_layout, layer_index, layer_pooling, token_pooling):
|
|
if cached is None:
|
|
raise gr.Error('先提取特征或导入缓存。')
|
|
layout = output_layout(cached['output_type']) if tensor_layout == 'auto' else tensor_layout
|
|
result = feature_tools(cached['features'], operation, int(reference_index), groups.splitlines() if groups.strip() else None,
|
|
tensor_layout=layout, layer_index=int(layer_index), layer_pooling=layer_pooling, token_pooling=token_pooling)
|
|
return json.dumps(result, ensure_ascii=False, indent=2, allow_nan=False)
|
|
|
|
|
|
ANALYSIS_OPTION_NAMES = (
|
|
'metric', 'normalize', 'cluster_method', 'n_clusters', 'top_k', 'reference_index',
|
|
'tensor_layout', 'layer_index', 'layer_pooling', 'token_pooling', 'labels', 'seed',
|
|
'perplexity', 'dbscan_eps', 'dbscan_min_samples', 'max_dimensions', 'heatmap_order',
|
|
'grid_width', 'width', 'height',
|
|
)
|
|
|
|
|
|
def build_analysis_tab():
|
|
"""Call within a Blocks/Tabs context; no ComfyUI or WebUI dependencies."""
|
|
with gr.TabItem('Features & Analysis'):
|
|
state = gr.State(value=None)
|
|
gr.Markdown('### 批量特征与分析\n先提取一次特征，再选择图表反复制图。也可以导入 `.npz` 缓存，无需加载模型。')
|
|
with gr.Row():
|
|
with gr.Column():
|
|
inputs = gr.File(label='图片批次（顺序以缓存 labels 为准）', file_count='multiple', file_types=['image'])
|
|
available = get_available_models()
|
|
model = gr.Dropdown(choices=available, value=(available or [None])[0], label='Model Folder')
|
|
refresh = gr.Button('刷新模型列表')
|
|
device = gr.Dropdown(['cuda', 'cpu'], value='cuda', label='Device')
|
|
output_type = gr.Dropdown(list(FEATURE_OUTPUTS), value='default', label='Feature Output')
|
|
layers = gr.Textbox(value='-1', label='Intermediate layers（逗号分隔 block 编号）')
|
|
norm = gr.Checkbox(value=True, label='Intermediate LayerNorm')
|
|
batch = gr.Number(value=2, precision=0, label='Inference batch size')
|
|
extract = gr.Button('提取并缓存特征')
|
|
with gr.Column():
|
|
cache_input = gr.File(label='导入 features.npz', file_types=['.npz'])
|
|
load = gr.Button('导入缓存（不推理）')
|
|
status = gr.Textbox(label='缓存状态', interactive=False)
|
|
cache_download = gr.File(label='下载特征缓存')
|
|
chart = gr.Dropdown(list(CHART_TYPES), value='relationship_graph', label='Chart Type')
|
|
with gr.Accordion('分析参数', open=False):
|
|
with gr.Row():
|
|
metric = gr.Dropdown(['cosine', 'euclidean', 'manhattan'], value='cosine', label='Distance metric')
|
|
normalize = gr.Checkbox(value=True, label='L2 normalize')
|
|
cluster = gr.Dropdown(['kmeans', 'agglomerative', 'dbscan', 'none'], value='kmeans', label='Clustering')
|
|
clusters = gr.Number(value=3, precision=0, label='Cluster count')
|
|
with gr.Row():
|
|
top_k = gr.Number(value=2, precision=0, label='Neighbor count')
|
|
reference = gr.Number(value=0, precision=0, label='Query image index (0-based)')
|
|
layout = gr.Dropdown(list(TENSOR_LAYOUTS), value='auto', label='Tensor layout')
|
|
layer_index = gr.Number(value=-1, precision=0, label='Input layer position')
|
|
with gr.Row():
|
|
layer_pooling = gr.Dropdown(['selected', 'mean'], value='selected', label='Layer pooling')
|
|
token_pooling = gr.Dropdown(['mean', 'flatten'], value='mean', label='Token pooling')
|
|
seed = gr.Number(value=42, precision=0, label='Seed')
|
|
perplexity = gr.Number(value=5.0, label='t-SNE perplexity')
|
|
with gr.Row():
|
|
eps = gr.Number(value=.35, label='DBSCAN eps')
|
|
min_samples = gr.Number(value=2, precision=0, label='DBSCAN min samples')
|
|
dimensions = gr.Number(value=32, precision=0, label='Max dimensions')
|
|
order = gr.Dropdown(['cluster', 'input'], value='cluster', label='Heatmap order')
|
|
with gr.Row():
|
|
grid_width = gr.Number(value=0, precision=0, label='Patch grid columns (0=auto)')
|
|
width = gr.Number(value=1400, precision=0, label='Width (px)')
|
|
height = gr.Number(value=1000, precision=0, label='Height (px)')
|
|
labels = gr.Textbox(value='', lines=3, label='Labels（每行一个，留空使用文件名）')
|
|
plot = gr.Button('从缓存生成图表')
|
|
visualization = gr.Image(label='Visualization', interactive=False)
|
|
report = gr.Textbox(label='Analysis JSON', lines=12, interactive=False)
|
|
downloads = gr.File(label='下载 PNG / JSON / 距离矩阵 CSV', file_count='multiple')
|
|
with gr.Accordion('共同特征 / 相似度 / 分组比较', open=False):
|
|
operation = gr.Dropdown(['common_features', 'similarity', 'compare_groups'], value='common_features', label='Operation')
|
|
groups = gr.Textbox(value='', lines=4, label='Group names（分组比较时，每张图片一个组名）')
|
|
run_tools = gr.Button('分析缓存特征')
|
|
tools_result = gr.Textbox(label='Feature result JSON', lines=8, interactive=False)
|
|
controls = [metric, normalize, cluster, clusters, top_k, reference, layout, layer_index,
|
|
layer_pooling, token_pooling, labels, seed, perplexity, eps, min_samples, dimensions,
|
|
order, grid_width, width, height]
|
|
extract.click(extract_uploaded, [inputs, model, device, output_type, layers, norm, batch], [state, status, cache_download])
|
|
load.click(import_uploaded_cache, [cache_input], [state, status])
|
|
plot.click(plot_uploaded_cache, [state, chart, *controls], [visualization, report, downloads])
|
|
run_tools.click(tools_uploaded_cache, [state, operation, reference, groups, layout, layer_index, layer_pooling, token_pooling], [tools_result])
|
|
refresh.click(lambda: gr.update(choices=get_available_models(), value=(get_available_models() or [None])[0]), [], [model])
|