primocode 10.0.0 → 10.0.9

package/studio/rastrear.py

@@ -0,0 +1,399 @@
+#!/usr/bin/env python3
+"""rastrear.py — onde está a pessoa, a cabeça e as mãos em CADA quadro.
+
+    "Tira o fundo e separa em camadas… põe o logo na minha mão… quando eu
+     juntar as mãos…" (30/09/2026)
+
+Tudo o que o editor faz "seguindo" alguém nasce aqui, numa passada só pelo
+vídeo, e fica guardado na pasta do projeto. Daí em diante as operações são
+conta sobre números (lib/video/rastro.js) — rápidas e testáveis.
+
+── O QUE SAI ────────────────────────────────────────────────────────────────
+  · matte-<asset>.mkv   a máscara da pessoa, cinza 0–255, um quadro por quadro
+                        da origem (FFV1: sem perda, e todo quadro é chave, então
+                        o -ss do ffmpeg cai no quadro exato — a máscara não
+                        escorrega um quadro em relação ao vídeo);
+  · fundo-<asset>.png   a "placa limpa": o fundo SEM a pessoa, montado com os
+                        pixels que a pessoa não cobriu ao longo do vídeo (e o
+                        buraco que sobra, preenchido por inpaint). É a camada
+                        de trás do "separa em camadas";
+  · rastro-<asset>.json por quadro: o rosto (caixa), o topo da cabeça (da
+                        máscara, não do rosto — o cabelo conta) e até duas
+                        mãos (centro da palma, tamanho, abertura, lado).
+
+── POR QUE A MÁSCARA PASSA POR UM FILTRO GUIADO ─────────────────────────────
+O segmentador do MediaPipe pensa em 256×256: ampliada para 1080×1920, a borda
+vira uma rampa larga e borrada — o "adesivo" que denuncia o recorte. O filtro
+guiado (He et al., 2010) usa o PRÓPRIO quadro como guia: a borda da máscara
+gruda na borda de verdade (ombro, cabelo) sem modelo nenhum a mais. A média
+com o quadro anterior tira o tremor de um quadro para o outro.
+
+── EM PARTES, EM PARALELO ───────────────────────────────────────────────────
+Cada processo analisa um trecho fixo de quadros (lib/video/rastrear.js
+reparte e roda vários ao mesmo tempo); o `juntar` costura as partes na
+ordem: as máscaras por concat sem recodificar, os quadros em fila, as somas
+da placa somadas.
+
+Uso (pelo uv, como o pessoa.py e o rostos.py — o MediaPipe só existe até o
+Python 3.12):
+    python rastrear.py analisar --entrada v.mp4 --pasta <midia> --id a1 [--de 96 --ate 192 --aquecer 12]
+    python rastrear.py juntar --pasta <midia> --id a1
+    python rastrear.py conferir
+"""
+
+import argparse
+import json
+import os
+import shutil
+import subprocess
+import sys
+import urllib.request
+
+MODELO_URL = ("https://storage.googleapis.com/mediapipe-models/image_segmenter/"
+              "selfie_multiclass_256x256/float32/latest/selfie_multiclass_256x256.tflite")
+PONTAS = (4, 8, 12, 16, 20)
+PALMA = (0, 5, 9, 13, 17)
+
+
+def achar_ffmpeg():
+    achado = shutil.which("ffmpeg")
+    if achado:
+        return achado
+    for c in ("/opt/homebrew/bin/ffmpeg", "/usr/local/bin/ffmpeg", "/usr/bin/ffmpeg"):
+        if os.path.exists(c):
+            return c
+    return None
+
+
+def modelo_multiclasse():
+    """O segmentador de 6 classes (fundo, cabelo, pele, roupa…) recorta o
+    cabelo bem melhor que o selfie_segmentation. Baixado uma vez; sem rede,
+    fica o antigo, que vem dentro do pacote."""
+    pasta = os.path.join(os.path.expanduser("~"), ".primocode", "studio", "modelos")
+    arq = os.path.join(pasta, "selfie_multiclass_256x256.tflite")
+    if os.path.exists(arq) and os.path.getsize(arq) > 100000:
+        return arq
+    try:
+        os.makedirs(pasta, exist_ok=True)
+        tmp = arq + ".tmp"
+        urllib.request.urlretrieve(MODELO_URL, tmp)
+        os.replace(tmp, arq)
+        return arq
+    except Exception:
+        return None
+
+
+class Segmentador:
+    def __init__(self, mp):
+        self.mp = mp
+        self.tipo = "selfie"
+        arq = modelo_multiclasse()
+        if arq:
+            try:
+                v = mp.tasks.vision
+                op = v.ImageSegmenterOptions(
+                    base_options=mp.tasks.BaseOptions(model_asset_path=arq),
+                    running_mode=v.RunningMode.VIDEO, output_confidence_masks=True,
+                    output_category_mask=False)
+                self.seg = v.ImageSegmenter.create_from_options(op)
+                self.tipo = "multiclasse"
+                return
+            except Exception:
+                pass
+        self.seg = mp.solutions.selfie_segmentation.SelfieSegmentation(model_selection=1)
+
+    def mascara(self, rgb, ms):
+        import numpy as np
+        if self.tipo == "multiclasse":
+            img = self.mp.Image(image_format=self.mp.ImageFormat.SRGB, data=rgb)
+            r = self.seg.segment_for_video(img, int(ms))
+            return 1.0 - np.asarray(r.confidence_masks[0].numpy_view(), dtype=np.float32)
+        return self.seg.process(rgb).segmentation_mask.astype(np.float32)
+
+    def fechar(self):
+        try:
+            self.seg.close()
+        except Exception:
+            pass
+
+
+def guiado(guia, p, r=8, eps=1e-2, fator=2):
+    """Filtro guiado rápido: a/b calculados em meia resolução, aplicados na cheia."""
+    import cv2
+    h, w = guia.shape
+    gi = cv2.resize(guia, (w // fator, h // fator), interpolation=cv2.INTER_AREA)
+    pi = cv2.resize(p, (w // fator, h // fator), interpolation=cv2.INTER_AREA)
+    rr = max(1, r // fator)
+    box = lambda x: cv2.boxFilter(x, -1, (2 * rr + 1, 2 * rr + 1))
+    mi, mp_ = box(gi), box(pi)
+    cov = box(gi * pi) - mi * mp_
+    var = box(gi * gi) - mi * mi
+    a = cov / (var + eps)
+    b = mp_ - a * mi
+    a = cv2.resize(box(a), (w, h), interpolation=cv2.INTER_LINEAR)
+    b = cv2.resize(box(b), (w, h), interpolation=cv2.INTER_LINEAR)
+    return a * guia + b
+
+
+def maos_do_quadro(res, w, h, pl=0, pt=0, W=None, H=None):
+    """As mãos em coordenadas do quadro ORIGINAL (0–1), descontada a borda."""
+    import math
+    W = W or w
+    H = H or h
+    out = []
+    if not res.multi_hand_landmarks:
+        return out
+    pares = list(zip(res.multi_hand_landmarks, res.multi_handedness or [None] * 9))
+    pares.sort(key=lambda x: -(x[1].classification[0].score if x[1] else 0))
+    for lm, mao in pares[:2]:
+        pts = [((p.x * W - pl) / w, (p.y * H - pt) / h) for p in lm.landmark]
+        cx = sum(pts[i][0] for i in PALMA) / len(PALMA)
+        cy = sum(pts[i][1] for i in PALMA) / len(PALMA)
+        # Tamanho da palma em fração da ALTURA (o mesmo para x e y em pixel).
+        palma = max(1e-6, math.hypot((pts[9][0] - pts[0][0]) * w, (pts[9][1] - pts[0][1]) * h))
+        abre = sum(math.hypot((pts[i][0] - pts[0][0]) * w, (pts[i][1] - pts[0][1]) * h) for i in PONTAS[1:]) / 4 / palma
+        rot = mao.classification[0] if mao else None
+        out.append([round(cx, 4), round(cy, 4), round(palma / h, 4), round(abre, 3),
+                    (rot.label[0] if rot else "?"), round(rot.score if rot else 0, 2)])
+    return out
+
+
+def medidas(entrada):
+    import cv2
+    cap = cv2.VideoCapture(entrada)
+    if not cap.isOpened():
+        return None
+    r = (cap.get(cv2.CAP_PROP_FPS) or 30, int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
+    cap.release()
+    return r
+
+
+def quadros_de(ff, entrada, w, h, fps, inicio, quantos):
+    """Os quadros [inicio, inicio+quantos) pelo ffmpeg, um a um (BGR).
+
+    O -ss meio quadro ANTES do primeiro: o ffmpeg decodifica desde o quadro
+    chave e descarta o que vem antes do tempo pedido, então o primeiro que
+    sai é exatamente o quadro `inicio` — a parte 3 de um rastreio em
+    paralelo começa no mesmo quadro em qualquer máquina."""
+    import numpy as np
+    args = [ff, "-hide_banner", "-loglevel", "error"]
+    if inicio > 0:
+        args += ["-ss", f"{(inicio - 0.5) / fps:.6f}"]
+    args += ["-i", entrada]
+    if quantos is not None:
+        args += ["-frames:v", str(quantos)]
+    args += ["-fps_mode", "passthrough", "-f", "rawvideo", "-pix_fmt", "bgr24", "-"]
+    p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
+    tam = w * h * 3
+    try:
+        while True:
+            b = p.stdout.read(tam)
+            if len(b) < tam:
+                break
+            yield np.frombuffer(b, np.uint8).reshape(h, w, 3)
+    finally:
+        p.stdout.close()
+        p.wait()
+
+
+def analisar(entrada, pasta, ident, de=0, ate=None, aquecer=0):
+    """Analisa os quadros [de, ate) (ate=None: até o fim) e grava a PARTE.
+
+    `aquecer` quadros antes de `de` passam pelos modelos e não são gravados:
+    o rastreador de mão e a média da máscara chegam à emenda já "quentes",
+    então a parte começa como se o vídeo tivesse sido lido desde o início.
+    As emendas ficam em quadros FIXOS (o Node corta de 96 em 96), não no
+    número de processos: o resultado é o mesmo com 1 ou com 12 deles."""
+    import cv2
+    import numpy as np
+    import mediapipe as mp
+
+    med = medidas(entrada)
+    if not med:
+        return {"ok": False, "error": f"não consegui abrir {entrada}"}
+    fps, w, h = med
+    ff = achar_ffmpeg()
+    if not ff:
+        return {"ok": False, "error": "falta o ffmpeg"}
+    dir_partes = os.path.join(pasta, f"partes-{ident}")
+    os.makedirs(dir_partes, exist_ok=True)
+    base = os.path.join(dir_partes, f"{de:06d}")
+    tmp_matte = base + ".tmp.mkv"
+    esc = subprocess.Popen([ff, "-hide_banner", "-loglevel", "error", "-y", "-f", "rawvideo", "-pix_fmt", "gray",
+                            "-s", f"{w}x{h}", "-r", str(fps), "-i", "-", "-c:v", "ffv1", "-level", "3",
+                            "-g", "1", "-pix_fmt", "gray", tmp_matte], stdin=subprocess.PIPE, stderr=subprocess.PIPE)
+
+    seg = Segmentador(mp)
+    maos = mp.solutions.hands.Hands(static_image_mode=False, max_num_hands=2, model_complexity=1,
+                                    min_detection_confidence=0.35, min_tracking_confidence=0.35)
+    rosto = mp.solutions.face_detection.FaceDetection(model_selection=1, min_detection_confidence=0.5)
+
+    # A placa limpa em meia resolução: soma dos pixels de fundo e quantas vezes cada um foi visto.
+    pw, ph = w // 2, h // 2
+    soma = np.zeros((ph, pw, 3), np.float64)
+    conta = np.zeros((ph, pw), np.float64)
+    anterior = None
+    quadros = []
+    inicio = max(0, de - aquecer)
+    quantos = None if ate is None else ate - inicio
+    k = inicio
+    try:
+        for bgr in quadros_de(ff, entrada, w, h, fps, inicio, quantos):
+            grava = k >= de
+            rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
+            m = seg.mascara(rgb, k * 1000.0 / fps)
+            if m.shape != (h, w):
+                m = cv2.resize(m, (w, h), interpolation=cv2.INTER_LINEAR)
+            guia = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY).astype(np.float32) / 255.0
+            m = guiado(guia, m, eps=1e-2)
+            # A curva: o segmentador deixa ~0,07 no fundo e uma rampa larga no
+            # ombro; sem ela o fundo velho aparece como fantasma atrás da pessoa.
+            m = np.clip((m - 0.3) / 0.4, 0, 1)
+            if anterior is not None:
+                m = 0.7 * m + 0.3 * anterior
+            anterior = m
+            # A mão que sai pela borda de baixo (o normal em quem fala para a
+            # câmera) quase nunca é achada: o detector de palma quer a mão
+            # inteira. Esticar a borda (replicada) dá a ele onde "terminar" a
+            # mão — medido no vídeo de teste: de 119 para 247 quadros com mão.
+            pl, pb = int(w * 0.12), int(h * 0.25)
+            borda = cv2.copyMakeBorder(rgb, 0, pb, pl, pl, cv2.BORDER_REPLICATE)
+            hs = maos_do_quadro(maos.process(borda), w, h, pl, 0, w + 2 * pl, h + pb)
+            k += 1
+            if not grava:
+                continue
+            esc.stdin.write((m * 255).astype(np.uint8).tobytes())
+
+            # Fundo: só o que está LONGE da pessoa (a máscara dilatada), para a borda não sujar a placa.
+            peq = cv2.resize(m, (pw, ph), interpolation=cv2.INTER_AREA)
+            livre = (cv2.dilate((peq > 0.15).astype(np.uint8), np.ones((15, 15), np.uint8)) == 0)
+            soma[livre] += cv2.resize(bgr, (pw, ph), interpolation=cv2.INTER_AREA)[livre]
+            conta[livre] += 1
+
+            q = {}
+            rf = rosto.process(rgb)
+            if rf.detections:
+                d = max(rf.detections, key=lambda x: x.score[0])
+                b = d.location_data.relative_bounding_box
+                q["f"] = [round(b.xmin, 4), round(b.ymin, 4), round(b.xmin + b.width, 4), round(b.ymin + b.height, 4), round(d.score[0], 2)]
+                # O topo da cabeça: a primeira linha com pessoa na faixa do rosto (o cabelo entra).
+                x0 = int(max(0, b.xmin + b.width * 0.3) * w)
+                x1 = int(min(1, b.xmin + b.width * 0.7) * w)
+                y1 = int(min(1, b.ymin + b.height * 0.5) * h)
+                if x1 > x0 and y1 > 0:
+                    faixa = m[:y1, x0:x1].mean(axis=1)
+                    linhas = np.where(faixa > 0.5)[0]
+                    if len(linhas):
+                        q["t"] = round(float(linhas[0]) / h, 4)
+            if hs:
+                q["h"] = hs
+            quadros.append(q)
+    finally:
+        try:
+            esc.stdin.close()
+        except Exception:
+            pass
+        seg.fechar()
+        maos.close()
+        rosto.close()
+    esc.wait()
+    if esc.returncode != 0:
+        return {"ok": False, "error": "o ffmpeg recusou a máscara: " + (esc.stderr.read() or b"").decode("utf-8", "ignore")[-300:]}
+    os.replace(tmp_matte, base + ".mkv")
+    np.savez(base + ".npz", soma=soma, conta=conta)
+    with open(base + ".json", "w") as f:
+        json.dump({"de": de, "fps": fps, "w": w, "h": h, "segmentador": seg.tipo, "quadros": quadros}, f, separators=(",", ":"))
+    return {"ok": True, "parte": de, "quadros": len(quadros)}
+
+
+def juntar(pasta, ident):
+    """As partes, em ordem de quadro, viram o rastro, a máscara e a placa."""
+    import cv2
+    import numpy as np
+    dir_partes = os.path.join(pasta, f"partes-{ident}")
+    nomes = sorted(n[:-5] for n in os.listdir(dir_partes) if n.endswith(".json"))
+    if not nomes:
+        return {"ok": False, "error": "nenhuma parte para juntar"}
+    quadros, soma, conta, meta = [], None, None, None
+    usadas = []
+    for nome in nomes:
+        with open(os.path.join(dir_partes, nome + ".json")) as f:
+            parte = json.load(f)
+        # O contêiner (.mov do celular) promete mais quadros do que o vídeo
+        # tem: a duração é a do áudio. A parte do fim sai vazia, e a anterior
+        # curta — é o fim do vídeo, não um buraco.
+        if not parte["quadros"]:
+            continue
+        if len(quadros) != parte["de"]:
+            return {"ok": False, "error": f"parte {nome} começa em {parte['de']}, esperava {len(quadros)}"}
+        quadros += parte["quadros"]
+        usadas.append(nome)
+        meta = meta or parte
+        z = np.load(os.path.join(dir_partes, nome + ".npz"))
+        soma = z["soma"] if soma is None else soma + z["soma"]
+        conta = z["conta"] if conta is None else conta + z["conta"]
+    w, h, fps = meta["w"], meta["h"], meta["fps"]
+    ff = achar_ffmpeg()
+    lista = os.path.join(dir_partes, "lista.txt")
+    with open(lista, "w") as f:
+        for nome in usadas:
+            f.write(f"file '{os.path.join(dir_partes, nome + '.mkv')}'\n")
+    matte = os.path.join(pasta, f"matte-{ident}.mkv")
+    r = subprocess.run([ff, "-hide_banner", "-loglevel", "error", "-y", "-f", "concat", "-safe", "0", "-i", lista,
+                        "-c", "copy", matte + ".tmp.mkv"], capture_output=True)
+    if r.returncode != 0:
+        return {"ok": False, "error": "não juntei as máscaras: " + r.stderr.decode("utf-8", "ignore")[-300:]}
+    os.replace(matte + ".tmp.mkv", matte)
+
+    ph, pw = conta.shape
+    placa = np.zeros((ph, pw, 3), np.uint8)
+    visto = conta > 0
+    placa[visto] = (soma[visto] / conta[visto][:, None]).astype(np.uint8)
+    buraco = (~visto).astype(np.uint8) * 255
+    if buraco.any():
+        placa = cv2.inpaint(placa, buraco, 9, cv2.INPAINT_TELEA)
+    fundo = os.path.join(pasta, f"fundo-{ident}.png")
+    cv2.imwrite(fundo, cv2.resize(placa, (w, h), interpolation=cv2.INTER_CUBIC))
+
+    dados = {"versao": 2, "fps": round(fps, 4), "w": w, "h": h, "n": len(quadros), "segmentador": meta["segmentador"],
+             "partes": len(nomes), "buraco": round(float((~visto).mean()), 3), "quadros": quadros}
+    arq = os.path.join(pasta, f"rastro-{ident}.json")
+    with open(arq, "w") as f:
+        json.dump(dados, f, separators=(",", ":"))
+    shutil.rmtree(dir_partes, ignore_errors=True)
+    return {"ok": True, "rastro": arq, "matte": matte, "fundo": fundo, "quadros": len(quadros), "segmentador": meta["segmentador"],
+            "comRosto": sum(1 for q in quadros if "f" in q), "comMao": sum(1 for q in quadros if "h" in q)}
+
+
+def main():
+    p = argparse.ArgumentParser(description="Rastreio de pessoa, rosto e mãos (PrimoCode)")
+    sub = p.add_subparsers(dest="comando", required=True)
+    sub.add_parser("conferir")
+    a = sub.add_parser("analisar")
+    a.add_argument("--entrada", required=True)
+    a.add_argument("--pasta", required=True)
+    a.add_argument("--id", required=True)
+    a.add_argument("--de", type=int, default=0)
+    a.add_argument("--ate", type=int, default=None)
+    a.add_argument("--aquecer", type=int, default=0)
+    j = sub.add_parser("juntar")
+    j.add_argument("--pasta", required=True)
+    j.add_argument("--id", required=True)
+    args = p.parse_args()
+    if args.comando == "conferir":
+        try:
+            import cv2, mediapipe, numpy  # noqa: F401
+            print(json.dumps({"ok": True}))
+        except Exception as e:
+            print(json.dumps({"ok": False, "error": str(e)}))
+        return
+    if args.comando == "juntar":
+        r = juntar(args.pasta, args.id)
+    else:
+        r = analisar(args.entrada, args.pasta, args.id, args.de, args.ate, args.aquecer)
+    print(json.dumps(r, ensure_ascii=False))
+    sys.exit(0 if r.get("ok") else 1)
+
+
+if __name__ == "__main__":
+    main()