import cv2, numpy as np, json
from PIL import Image, ImageDraw, ImageFont
F='/usr/share/fonts/opentype/noto/NotoSansCJK-Bold.ttc'
def font(s): return ImageFont.truetype(F,s,index=0)  # JP
base=Image.open('crop.png').convert('RGB'); W,H=base.size
M={n:np.load(f'm_{n}.npy') for n in ['arm','red','blue']}
# arm: fill holes
a=(M['arm']*255).astype(np.uint8); cs,_=cv2.findContours(a,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
a2=np.zeros_like(a); cv2.drawContours(a2,cs,-1,255,-1); M['arm']=a2>0
def bbox(m): ys,xs=np.where(m); return [int(xs.min()),int(ys.min()),int(xs.max()),int(ys.max())]
B={n:bbox(m) for n,m in M.items()}
NAME={'arm':'アーム','red':'赤いブロック','blue':'青いブロック'}
KP=[(470,95),(470,265),(462,345),(458,405),(395,568),(508,568)]
def chip(d,xy,text,size=22,fg='white',bg=(20,20,20)):
    f=font(size); x,y=xy; tw=d.textlength(text,font=f); d.rectangle([x,y,x+tw+16,y+size+12],fill=bg); d.text((x+8,y+3),text,font=f,fill=fg)
def dim(im,k=0.55): return Image.blend(im,Image.new('RGB',im.size,(0,0,0)),1-k)
out={}
# 分類
im=base.copy(); d=ImageDraw.Draw(im); chip(d,(24,24),'机の上のロボットアーム',38); out['cls']=im
# 検出
im=base.copy(); d=ImageDraw.Draw(im)
for n,(x0,y0,x1,y1) in B.items():
    p=8; d.rectangle([x0-p,y0-p,x1+p,y1+p],outline='white',width=5); chip(d,(x0-p,y0-p-52),NAME[n],30)
out['det']=im
# セグメンテーション：暗くして、物の形だけ白く塗る
im=np.array(dim(base,0.45)).astype(float)
for n,v in [('arm',255),('red',255),('blue',255)]:
    m=M[n]; im[m]=im[m]*0.25+np.array([v,v,v])*0.75
im=Image.fromarray(im.astype(np.uint8)); d=ImageDraw.Draw(im)
for n in M:
    cs,_=cv2.findContours((M[n]*255).astype(np.uint8),cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_NONE)
    for c in cs: d.line([tuple(p[0]) for p in c]+[tuple(c[0][0])],fill='white',width=3)
out['seg']=im
# 姿勢推定
im=dim(base,0.6); d=ImageDraw.Draw(im)
d.line(KP[:4],fill='white',width=6); d.line([KP[3],KP[4]],fill='white',width=6); d.line([KP[3],KP[5]],fill='white',width=6)
for x,y in KP: d.ellipse([x-13,y-13,x+13,y+13],fill='white',outline=(20,20,20),width=4)
out['pose']=im
# 追跡（説明用：手の位置の通り道）
im=dim(base,0.6); d=ImageDraw.Draw(im)
trail=[(150,250),(230,300),(320,350),(400,390),(458,405)]
for i in range(len(trail)-1): d.line([trail[i],trail[i+1]],fill='white',width=5)
for x,y in trail[:-1]: d.ellipse([x-7,y-7,x+7,y+7],fill='white')
x0,y0,x1,y1=B['arm']; chip(d,(trail[-1][0]+20,trail[-1][1]-60),'ID 1  アームの手',22)
for n in ['red','blue']:
    x0,y0,x1,y1=B[n]; d.rectangle([x0-8,y0-8,x1+8,y1+8],outline='white',width=4); chip(d,(x0-8,y0-48),'ID 2' if n=='red' else 'ID 3',22)
out['trk']=im
# 深度（Depth Anything V2 Small の出力。近いほど白）
dp=np.load('depth.npy'); dn=(dp-dp.min())/(dp.max()-dp.min())
im=Image.fromarray((dn*255).astype(np.uint8)).convert('RGB'); out['dep']=im
for k,v in out.items(): v.save(f'A_{k}.png')
json.dump({'bbox':B,'kp':KP,'area':{n:int(M[n].sum()) for n in M}},open('meta.json','w'),ensure_ascii=False)
# データ案用：セグの粗い格子（24x16）
g=np.zeros((16,24),int)
for n,v in [('arm',1),('red',2),('blue',3)]:
    s=cv2.resize(M[n].astype(np.uint8),(24,16),interpolation=cv2.INTER_AREA); g[s>0]=v
json.dump(g.tolist(),open('grid.json','w'))
dd=cv2.resize(dn,(8,5),interpolation=cv2.INTER_AREA); json.dump(np.round(dd,2).tolist(),open('dgrid.json','w'))
# 輪郭の点の数（手で点を打つ数の目安）
pts={}
for n in M:
    cs,_=cv2.findContours((M[n]*255).astype(np.uint8),cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_NONE)
    c=max(cs,key=cv2.contourArea); pts[n]=len(cv2.approxPolyDP(c,2.0,True))
print(B,pts)
json.dump(pts,open('pts.json','w'))