#!/usr/bin/env python3
"""看图识象 Shixiang v3 —— 看图认物象（山水花鸟 121 类，其中垂钓低可信，月在照片上不可用、火、塔不可用）＋画种（6 类）。独立推理参考实现（PyTorch）。
结构：图像编码器（Mitsua Japanese CLIP ViT-B/16 经人判古画切块微调，CC BY-SA 4.0）→ 多视角特征（整图＋2×2 四格；可选再加 5 个滑窗，共 10 视角）
      → 分类头（多示例：每视角打分取最大，顺带给位置）。
文件：encoder.safetensors（图像编码器＋投影，fp16）、head.pt（分类头）、zs_text.npz（19 类零样本文字向量，塔等用）、zs_rules.json（零样本补标规则）、
      thresholds.json（照片门槛，th＝5 视角、th_slide＝加滑窗）、paint_thresholds.json（画上门槛，th_paint／th_paint_slide）、vocab.json。
用法：python3 shixiang.py [--slide] 图片.jpg [更多图片]   或   from shixiang import Shixiang; m = Shixiang('.'); m.tag(PIL.Image, slide=True)"""
import os, json, numpy as np, torch, torch.nn as nn, torch.nn.functional as F
from PIL import Image

VIEWS = ['整图', '左上', '右上', '左下', '右下', '上中', '左中', '中央', '右中', '下中']
SLIDE = [(0, 112), (112, 0), (112, 112), (112, 224), (224, 112)]   # 滑窗左上角 (y, x)，448 图上 224 窗、步长 112，与四格合成 3×3 九窗
# 零样本第二意见的规则（哪些词可单独补标、分类头最低把握、免检词、舟船合并）在 zs_rules.json，作为数据随模型发布。


def views_uint8(im: Image.Image, slide: bool = False) -> np.ndarray:
    """PIL → uint8 [5 或 10,224,224,3]：整图拉伸到 224（同官方 SiglipImageProcessor），四格＝拉伸到 448 再四分；slide 再加 5 个滑窗。"""
    im = im.convert('RGB')
    g = np.asarray(im.resize((224, 224), Image.BICUBIC)); q = np.asarray(im.resize((448, 448), Image.BICUBIC))
    v = [g, q[:224, :224], q[:224, 224:], q[224:, :224], q[224:, 224:]]
    if slide: v += [q[y:y + 224, x:x + 224] for y, x in SLIDE]
    return np.stack(v)


def to_tensor(v: np.ndarray) -> torch.Tensor:
    """uint8 [B,V,224,224,3] → float [B,V,3,224,224]，均值 0.5 方差 0.5。"""
    x = torch.from_numpy(np.ascontiguousarray(v)).permute(0, 1, 4, 2, 3).float()
    return (x / 255. - 0.5) / 0.5


class Head(nn.Module):
    """分数 = 缩放·余弦(视角特征, 物象文字向量+可学偏移) + 偏置 + MLP 残差；画种用整图视角。"""
    def __init__(s, V, d=768, h=1024, nk=6):
        super().__init__(); s.W = nn.Parameter(torch.zeros(V, d)); s.ls = nn.Parameter(torch.zeros(())); s.b = nn.Parameter(torch.zeros(V))
        s.f = nn.Sequential(nn.Linear(d, h), nn.GELU(), nn.Dropout(0.1)); s.obj = nn.Linear(h, V); s.kind = nn.Linear(d, nk)
    def forward(s, x):   # x: B,V,768 已 L2 归一
        return s.ls.exp() * (x @ F.normalize(s.W, dim=-1).T) + s.b + s.obj(s.f(x)), s.kind(x[:, 0])


class ShixiangNet(nn.Module):
    """[B,V,3,224,224] → feats [B,V,768]、logits [B,V,C]、kind [B,6]、zs [B,19]。V＝5 或 10，视角 0 必须是整图。"""
    def __init__(s, d, zs_T, zs_D):
        super().__init__()
        from transformers import CLIPVisionConfig, CLIPVisionModel
        cfg = CLIPVisionConfig(hidden_size=768, intermediate_size=3072, num_hidden_layers=12, num_attention_heads=12, image_size=224, patch_size=16,
                               hidden_act='gelu', layer_norm_eps=1e-6, projection_dim=768, attn_implementation='eager')
        s.vision = CLIPVisionModel(cfg).vision_model; s.proj = nn.Linear(768, 768, bias=False); s.head = Head(d['V'])
        s.register_buffer('zs_T', zs_T); s.register_buffer('zs_D', zs_D)
    def forward(s, x):
        B, V = x.shape[0], x.shape[1]; x = x.reshape(B * V, 3, 224, 224)
        f = s.proj(s.vision(pixel_values=x)[1]); f = F.normalize(f, dim=-1).reshape(B, V, -1)
        lo, kd = s.head(f); x0 = f[:, 0]
        zs = x0 @ s.zs_T.T - (x0 @ s.zs_D.T).mean(1, keepdim=True)
        return f, lo, kd, zs


class Shixiang:
    def __init__(s, d='.', device='cpu'):
        s.dir = d; s.device = device
        s.vocab = json.load(open(f'{d}/vocab.json', encoding='utf-8')); s.VOC = s.vocab['vocab']; s.KINDS = s.vocab['kinds']
        th = json.load(open(f'{d}/thresholds.json', encoding='utf-8'))['thresholds']
        s.th = np.array([th[c]['th'] if th[c]['th'] is not None else 2.0 for c in s.VOC], np.float32)
        s.th_slide = np.array([th[c].get('th_slide', th[c]['th']) if th[c]['th'] is not None else 2.0 for c in s.VOC], np.float32)
        s.conf = {c: th[c]['confidence'] for c in s.VOC}
        s.paint = json.load(open(f'{d}/paint_thresholds.json', encoding='utf-8'))
        z = np.load(f'{d}/zs_text.npz'); s.ZC = [str(c) for c in z['cls']]
        s.zr = json.load(open(f'{d}/zs_rules.json', encoding='utf-8')); assert s.zr['classes'] == s.ZC
        s.ZTH = np.array([s.zr['th'][c] for c in s.ZC], np.float32); s.MERGE = s.zr['merge']
        s.net = ShixiangNet(dict(V=len(s.VOC)), torch.from_numpy(z['T']), torch.from_numpy(z['D']))
        from safetensors.torch import load_file
        enc = load_file(f'{d}/encoder.safetensors')
        s.net.vision.load_state_dict({k[len('vision_model.'):]: v.float() for k, v in enc.items() if k.startswith('vision_model.')})
        s.net.proj.load_state_dict({'weight': enc['visual_projection.weight'].float()})
        ck = torch.load(f'{d}/head.pt', map_location='cpu'); s.net.head.load_state_dict(ck['state_dict']); assert ck['vocab'] == s.VOC
        s.net.eval().to(device); s.vi = {c: i for i, c in enumerate(s.VOC)}

    @torch.no_grad()
    def forward(s, ims, slide=False):
        x = to_tensor(np.stack([views_uint8(im, slide) for im in ims])).to(s.device)
        f, lo, kd, zs = s.net(x)
        P = torch.sigmoid(lo); sc, vw = P.max(1)
        return dict(feats=f.cpu().numpy(), scores=sc.cpu().numpy(), view=vw.cpu().numpy(), kind=torch.softmax(kd, -1).cpu().numpy(), zs=zs.cpu().numpy())

    def decide(s, scores, view, kind, zs, painting=None, slide=False):
        """一张图的分数 → 标签表。painting=None 时按画种自动判（判为国画才用画上门槛，见 paint_thresholds.json 的 _kinds）。slide=True 用加滑窗的门槛。"""
        kind_i = int(kind.argmax()); kname = s.KINDS[kind_i]
        if painting is None: painting = kname in s.paint.get('_kinds', ['国画'])   # 画上门槛只给国画（paint_thresholds.json 的 _kinds）
        th = (s.th_slide if slide else s.th).copy()
        if painting:   # 画上门槛：按类别词组覆盖（校准自人判画块验收集）
            for cat, r in s.paint.items():
                if cat.startswith('_'): continue
                t = r.get('th_paint_slide', r['th_paint']) if slide else r['th_paint']
                for w in r['words']:
                    if w in s.vi: th[s.vi[w]] = t
        if not painting and kname != '照片':   # 油画、版画、素描、水彩：个别类另设最低把握（paint_thresholds.json 的 _other_art_min）
            for w, mn in s.paint.get('_other_art_min', {}).items():
                if w in s.vi: th[s.vi[w]] = max(th[s.vi[w]], mn)
        tg = {}
        for j in np.where(scores >= th)[0]:
            w = s.MERGE.get(s.VOC[j], s.VOC[j]); sc = float(scores[j])
            if w not in tg or sc > tg[w]['score']: tg[w] = dict(word=w, score=round(sc, 3), where=VIEWS[int(view[j])], src='head', confidence=s.conf[s.VOC[j]])
        for k in np.where(zs >= s.ZTH)[0]:
            c = s.ZC[k]
            h = scores[s.vi[c]] if c in s.vi else 0.0; zr = s.zr
            if (c in zr['alone'] and (c in zr['alone_exempt'] or h >= zr['alone_min_head'])) or (c in zr['weak'] and h >= zr['weak_min_head']):
                w = s.MERGE.get(c, c)
                if w not in tg: tg[w] = dict(word=w, score=round(float(zs[k]), 3), where='整图', src='zero-shot', confidence='正常')
                else: tg[w]['src'] = 'both'
        tags = sorted(tg.values(), key=lambda t: -t['score'])
        return dict(kind=kname, kind_prob=round(float(kind[kind_i]), 3), painting=bool(painting), tags=tags)

    def tag(s, ims, painting=None, slide=False):
        one = isinstance(ims, Image.Image); ims = [ims] if one else list(ims)
        o = s.forward(ims, slide)
        res = [s.decide(o['scores'][i], o['view'][i], o['kind'][i], o['zs'][i], painting, slide) for i in range(len(ims))]
        return res[0] if one else res


if __name__ == '__main__':
    import sys
    args = [p for p in sys.argv[1:] if p != '--slide']; slide = '--slide' in sys.argv
    m = Shixiang(os.path.dirname(os.path.abspath(__file__)) or '.')
    for p in args:
        r = m.tag(Image.open(p), slide=slide); print(p, r['kind'], ' '.join(f"{t['word']}({t['score']:.2f},{t['where']})" for t in r['tags']))