"""Infer each example once, cache features, then exercise every plotting node mode."""
import argparse
import csv
import hashlib
import html
import json
from pathlib import Path
import re
import sys
import numpy as np
from PIL import Image, ImageDraw, ImageFont, ImageOps
import torch
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from feature_analysis import CHART_TYPES
from kaloscope_dinov3.preprocessing import prepare_rgb
from model_loading import load_model_bundle
from test_model_loading import load_nodes
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--input', default='../test')
parser.add_argument('--model-dir', default='../model')
parser.add_argument('--output', default='../outputs')
parser.add_argument('--device', default='cuda')
parser.add_argument('--batch-size', default=2, type=int)
parser.add_argument('--reuse-cache', action='store_true', help='Redraw charts from features.pt without loading the model')
args = parser.parse_args()
torch.set_num_threads(4)
source = Path(args.input).resolve()
output = Path(args.output).resolve()
output.mkdir(parents=True, exist_ok=True)
paths = sorted(path for path in source.iterdir() if path.suffix.lower() in {'.png','.jpg','.jpeg','.webp','.bmp'})
if not paths:
raise ValueError(f'No images in {source}')
if args.batch_size < 1:
raise ValueError('batch-size must be positive')
cache_path=output/'features.pt'
if args.reuse_cache:
cached=torch.load(cache_path,map_location='cpu',weights_only=True)
features,patch_features=cached['features'],cached['patch_tokens']
labels,records=cached['labels'],cached['source_images']
model_info=cached.get('model_info',{'model_type':'dinov3_vitb16','feature_source':'projector',
'input_size':512,'checkpoint':str(Path(args.model_dir).resolve()/'best.pt')})
inference_batches=cached.get('inference_batches',4)
previews=[]
for record in records:
with Image.open(record['path']) as original:
previews.append(ImageOps.pad(prepare_rgb(original),(240,200),color='white'))
thumbnails=torch.stack([torch.from_numpy(np.array(image)).float()/255 for image in previews])
print('Reusing cached features: zero new model forwards',flush=True)
else:
bundle = load_model_bundle(args.model_dir, device=args.device)
model = bundle['model']
model_info = {key: bundle[key] for key in ('model_type','feature_source','input_size','checkpoint')}
if not hasattr(model, 'extract_tensor'):
raise ValueError('This example suite requires a DINOv3 checkpoint for patch analysis')
vectors, patches, previews = [], [], []
labels, records = [], []
inference_batches = 0
with torch.inference_mode():
for start in range(0,len(paths),args.batch_size):
batch = []
for index, path in enumerate(paths[start:start+args.batch_size], start):
with Image.open(path) as original:
image = prepare_rgb(original)
batch.append(bundle['transform'](image))
previews.append(ImageOps.pad(image, (240,200), color='white'))
artist = re.search(r'_drawn_by_(.+?)__[0-9a-f]+$', path.stem)
label = f'{index+1:02d} ' + (artist.group(1) if artist else path.stem[:22])
labels.append(label)
with path.open('rb') as stream:
digest = hashlib.file_digest(stream, 'sha256').hexdigest()
records.append({'index': index, 'label': label, 'path': str(path), 'sha256': digest})
tokens = model.backbone.forward_features(torch.stack(batch).to(args.device))
pooled = model._pool(tokens['x_norm_clstoken'],tokens['x_norm_patchtokens'],model.pooling)
feature = model.projector(pooled) if model.feature_source=='projector' else pooled
vectors.append(feature.float().cpu())
patches.append(tokens['x_norm_patchtokens'].float().cpu())
inference_batches += 1
print(f'Inferred images {start+1}-{min(start+args.batch_size,len(paths))} once',flush=True)
features = torch.cat(vectors)
patch_features = torch.cat(patches)
thumbnails = torch.stack([torch.from_numpy(np.array(image)).float()/255 for image in previews])
torch.save({'features':features, 'patch_tokens':patch_features, 'labels':labels, 'source_images':records,'model_info':model_info,'inference_batches':inference_batches},output/'features.pt')
# Release GPU model; all charts below consume cached CPU tensors only.
del model, bundle
if torch.cuda.is_available():
torch.cuda.empty_cache()
nodes = load_nodes()
node = nodes.KaloscopeFeatureAnalysisNode()
generated = []
for chart in CHART_TYPES:
tensor = patch_features if chart=='patch_energy' else features
result, text, distances = node.analyze(tensor,chart_type=chart,labels=json.dumps(labels),images=thumbnails,
tensor_layout='tokens' if chart=='patch_energy' else 'vectors',n_clusters=3,top_k=3,
metric='cosine',normalize=True,width=1600,height=1100,perplexity=3.0)
report = json.loads(text)
report['source_images']=records
report['feature_source']='patch_tokens' if chart=='patch_energy' else model_info['feature_source']
array = (result[0].clamp(0,1).numpy()*255).round().astype(np.uint8)
Image.fromarray(array).save(output/f'{chart}.png',dpi=(120,120))
(output/f'{chart}.json').write_text(json.dumps(report,ensure_ascii=False,indent=2,allow_nan=False),encoding='utf-8')
if chart=='relationship_graph':
with (output/'distance_matrix.csv').open('w',encoding='utf-8-sig',newline='') as stream:
writer=csv.writer(stream)
writer.writerow(['image',*labels])
writer.writerows([label,*row.tolist()] for label,row in zip(labels,distances))
generated.append({'type':chart,'image':f'{chart}.png','report':f'{chart}.json'})
print(f'Saved {chart}.png',flush=True)
manifest={'input_directory':str(source),'checkpoint':model_info['checkpoint'],
'architecture':model_info['model_type'],'model_info':model_info,'feature_shape':list(features.shape),'patch_shape':list(patch_features.shape),
'source_images':records,'inference_batches':inference_batches,'image_forward_count':len(paths),
'analysis_model_forward_count':0,'this_run_model_forward_count':0 if args.reuse_cache else len(paths),'seed':42,'charts':generated}
(output/'manifest.json').write_text(json.dumps(manifest,ensure_ascii=False,indent=2),encoding='utf-8')
body=''.join(f'' for item in generated)
(output/'index.html').write_text('''Kaloscope 分析图示例
Kaloscope 特征分析示例
每张输入图片仅推理一次。18 类图表复用缓存特征；二维距离布局是近似结果，具体距离以矩阵和 JSON 为准。
来源与运行记录 · 距离矩阵 CSV
'''+body+'',encoding='utf-8')
contact = Image.new('RGB',(1600,math_rows(len(generated))*390),'#eef2f5')
drawer=ImageDraw.Draw(contact)
try:
font=ImageFont.truetype('C:/Windows/Fonts/arial.ttf',22)
except OSError:
font=ImageFont.load_default()
for index,item in enumerate(generated):
x=(index%3)*530+10; y=(index//3)*390
with Image.open(output/item['image']) as image:
preview=ImageOps.contain(image,(510,350))
contact.paste(preview,(x,y+32))
drawer.text((x,y+5),item['type'],fill='#243746',font=font)
contact.save(output/'overview.png')
print(f'{len(generated)} chart types written to {output}; analysis performed zero model forwards',flush=True)
def math_rows(count):
return (count+2)//3
if __name__=='__main__':
main()