primocode 10.16.0 → 10.17.0
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/LICENSES/edit-ia-pro.txt +32 -0
- package/LICENSES/fontes-OFL.txt +476 -0
- package/LICENSES/hyperframes-Apache-2.0.txt +205 -0
- package/LICENSES/permanent-marker-Apache-2.0.txt +204 -0
- package/LICENSES/selfie-multiclass-Apache-2.0.txt +227 -0
- package/LICENSES/three-MIT.txt +23 -0
- package/LICENSES/video-use-MIT.txt +24 -0
- package/LICENSES/woven-sfx-CC0.txt +16 -0
- package/LICENSES/yunet-MIT.txt +21 -0
- package/README.md +2 -2
- package/bin/primocode.js +11 -2
- package/editor/LICENSE.md +11 -0
- package/editor/README.md +321 -0
- package/editor/SKILL.md +307 -0
- package/editor/THIRD_PARTY_NOTICES.md +228 -0
- package/editor/assets/baixar.json +591 -0
- package/editor/assets/fonts/anton/OFL.txt +93 -0
- package/editor/assets/fonts/caveat/OFL.txt +93 -0
- package/editor/assets/fonts/figtree/OFL.txt +93 -0
- package/editor/assets/fonts/fonts.json +10 -0
- package/editor/assets/fonts/montserrat/OFL.txt +93 -0
- package/editor/assets/fonts/permanentmarker/LICENSE.txt +202 -0
- package/editor/assets/models/SELFIE_MULTICLASS_LICENSE.txt +227 -0
- package/editor/assets/models/YUNET_LICENSE.txt +21 -0
- package/editor/assets/sfx/licencas/bass_drop.txt +32 -0
- package/editor/assets/sfx/licencas/cash.txt +76 -0
- package/editor/assets/sfx/licencas/charge.txt +31 -0
- package/editor/assets/sfx/licencas/click.txt +46 -0
- package/editor/assets/sfx/licencas/coin.txt +50 -0
- package/editor/assets/sfx/licencas/count.txt +47 -0
- package/editor/assets/sfx/licencas/data.txt +31 -0
- package/editor/assets/sfx/licencas/error.txt +32 -0
- package/editor/assets/sfx/licencas/flash.txt +31 -0
- package/editor/assets/sfx/licencas/gears.txt +82 -0
- package/editor/assets/sfx/licencas/keyboard.txt +48 -0
- package/editor/assets/sfx/licencas/metal_hit.txt +62 -0
- package/editor/assets/sfx/licencas/metal_riser.txt +61 -0
- package/editor/assets/sfx/licencas/mouse_click.txt +50 -0
- package/editor/assets/sfx/licencas/notification.txt +54 -0
- package/editor/assets/sfx/licencas/pop.txt +38 -0
- package/editor/assets/sfx/licencas/rewind.txt +32 -0
- package/editor/assets/sfx/licencas/rise.txt +46 -0
- package/editor/assets/sfx/licencas/sinteticos.txt +21 -0
- package/editor/assets/sfx/licencas/snap.txt +50 -0
- package/editor/assets/sfx/licencas/success.txt +40 -0
- package/editor/assets/sfx/licencas/synth_hit.txt +32 -0
- package/editor/assets/sfx/licencas/tick.txt +14 -0
- package/editor/assets/sfx/licencas/whoosh.txt +47 -0
- package/editor/assets/sfx/licencas/wind.txt +31 -0
- package/editor/assets/sfx/sfx.json +832 -0
- package/editor/assets/sfx/sfx_sinteticos.json +317 -0
- package/editor/assets/sfx/sfx_woven.json +387 -0
- package/editor/assets/sfx/woven/LICENCA.md +16 -0
- package/editor/fichas/briefing-do-video.md +45 -0
- package/editor/fichas/conferencia-da-amostra.md +56 -0
- package/editor/fichas/matriz-de-ganchos.csv +3 -0
- package/editor/fichas/padrao.json +16 -0
- package/editor/fichas/planilha-de-cenas.csv +4 -0
- package/editor/fichas/planilha-de-versoes.csv +4 -0
- package/editor/fichas/plano-em-blocos.md +32 -0
- package/editor/fichas/registro-de-materiais.json +27 -0
- package/editor/helpers/assets_ocr/ocr_vision.swift +41 -0
- package/editor/helpers/bordas_precisas.py +272 -0
- package/editor/helpers/camera_viva.py +426 -0
- package/editor/helpers/captions.py +601 -0
- package/editor/helpers/common.py +308 -0
- package/editor/helpers/compose_final.py +3555 -0
- package/editor/helpers/conferir_corte.py +234 -0
- package/editor/helpers/conferir_dados.py +405 -0
- package/editor/helpers/contact_sheet.py +282 -0
- package/editor/helpers/corte_por_onda.py +433 -0
- package/editor/helpers/corte_seco.py +595 -0
- package/editor/helpers/denoise.py +971 -0
- package/editor/helpers/efeitos.py +408 -0
- package/editor/helpers/efeitos_gerar.py +159 -0
- package/editor/helpers/elevenlabs_api.py +442 -0
- package/editor/helpers/face_layout.py +570 -0
- package/editor/helpers/fetch_images.py +423 -0
- package/editor/helpers/ganho_por_trecho.py +214 -0
- package/editor/helpers/gen_image.py +312 -0
- package/editor/helpers/gen_music.py +316 -0
- package/editor/helpers/gen_sfx.py +148 -0
- package/editor/helpers/grade.py +307 -0
- package/editor/helpers/icloud.py +119 -0
- package/editor/helpers/make_sfx.py +940 -0
- package/editor/helpers/matte.py +1342 -0
- package/editor/helpers/medir_referencia.py +2232 -0
- package/editor/helpers/mixar_trilha.py +354 -0
- package/editor/helpers/montar_corte.py +295 -0
- package/editor/helpers/motion/__init__.py +13 -0
- package/editor/helpers/motion/engine.py +1915 -0
- package/editor/helpers/motion/registry.py +510 -0
- package/editor/helpers/motion/render.py +665 -0
- package/editor/helpers/motion/templates/barra_progresso.py +370 -0
- package/editor/helpers/motion/templates/carimbo.py +132 -0
- package/editor/helpers/motion/templates/checklist.py +271 -0
- package/editor/helpers/motion/templates/comentario_cta.py +380 -0
- package/editor/helpers/motion/templates/contador.py +222 -0
- package/editor/helpers/motion/templates/destaque.py +263 -0
- package/editor/helpers/motion/templates/etiqueta_preco.py +461 -0
- package/editor/helpers/motion/templates/grafico_barras.py +510 -0
- package/editor/helpers/motion/templates/grafico_linha.py +462 -0
- package/editor/helpers/motion/templates/notificacao.py +266 -0
- package/editor/helpers/motion/templates/palavra_impacto.py +410 -0
- package/editor/helpers/motion/templates/passos.py +321 -0
- package/editor/helpers/motion/templates/riscado.py +298 -0
- package/editor/helpers/motion/templates/selo_garantia.py +150 -0
- package/editor/helpers/motion_render.py +241 -0
- package/editor/helpers/motion_sugerir.py +961 -0
- package/editor/helpers/noise_check.py +134 -0
- package/editor/helpers/pack_transcripts.py +223 -0
- package/editor/helpers/palavras_na_onda.py +176 -0
- package/editor/helpers/plano_sfx.py +839 -0
- package/editor/helpers/portao.py +427 -0
- package/editor/helpers/preview_server.py +1177 -0
- package/editor/helpers/rastrear_rosto.py +298 -0
- package/editor/helpers/render_cut.py +1256 -0
- package/editor/helpers/respiros.py +89 -0
- package/editor/helpers/sincronizar_cameras.py +571 -0
- package/editor/helpers/speech_regions.py +423 -0
- package/editor/helpers/style_catalog.json +231 -0
- package/editor/helpers/style_samples.py +246 -0
- package/editor/helpers/tela_verde.py +2113 -0
- package/editor/helpers/text_render.py +2059 -0
- package/editor/helpers/texto_na_tela.py +366 -0
- package/editor/helpers/timeline_view.py +292 -0
- package/editor/helpers/tira_de_quadros.py +113 -0
- package/editor/helpers/transcribe.py +734 -0
- package/editor/helpers/transcribe_batch.py +129 -0
- package/editor/helpers/trilhas.py +1146 -0
- package/editor/helpers/trilhas_humores.json +66 -0
- package/editor/helpers/valores.py +187 -0
- package/editor/helpers/verify_cut.py +820 -0
- package/editor/helpers/voice_levels.py +233 -0
- package/editor/helpers/voz_check.py +386 -0
- package/editor/helpers/voz_limpa.py +81 -0
- package/editor/helpers/watch_preview.py +372 -0
- package/editor/modelo-projeto/LEIA.md +402 -0
- package/editor/modelo-projeto/assets/fontes/OFL-Caveat.txt +93 -0
- package/editor/modelo-projeto/assets/fontes/OFL-JetBrainsMono.txt +93 -0
- package/editor/modelo-projeto/assets/fontes/OFL-Montserrat.txt +93 -0
- package/editor/modelo-projeto/assets/fontes/larguras.json +1 -0
- package/editor/modelo-projeto/assets/three/BufferGeometryUtils.js +1438 -0
- package/editor/modelo-projeto/assets/three/GLTFLoader.js +4895 -0
- package/editor/modelo-projeto/assets/three/LICENSE +21 -0
- package/editor/modelo-projeto/assets/three/ORIGEM.txt +17 -0
- package/editor/modelo-projeto/assets/three/RoomEnvironment.js +185 -0
- package/editor/modelo-projeto/assets/three/three.core.min.js +6 -0
- package/editor/modelo-projeto/assets/three/three.module.min.js +6 -0
- package/editor/modelo-projeto/build.py +4828 -0
- package/editor/modelo-projeto/compositions/cenas/2d/k2d_antes_depois.html +199 -0
- package/editor/modelo-projeto/compositions/cenas/2d/k2d_arquivo.html +232 -0
- package/editor/modelo-projeto/compositions/cenas/2d/k2d_busca.html +315 -0
- package/editor/modelo-projeto/compositions/cenas/2d/k2d_comando.html +304 -0
- package/editor/modelo-projeto/compositions/cenas/2d/k2d_conversa.html +217 -0
- package/editor/modelo-projeto/compositions/cenas/2d/k2d_esquema.html +289 -0
- package/editor/modelo-projeto/compositions/cenas/2d/k2d_medidor.html +249 -0
- package/editor/modelo-projeto/compositions/cenas/2d/k2d_notas.html +233 -0
- package/editor/modelo-projeto/compositions/cenas/3d/k3d_caminho.html +197 -0
- package/editor/modelo-projeto/compositions/cenas/3d/k3d_explodida.html +152 -0
- package/editor/modelo-projeto/compositions/cenas/3d/k3d_tela.html +167 -0
- package/editor/modelo-projeto/compositions/cenas/3d/k3d_volume.html +124 -0
- package/editor/modelo-projeto/compositions/cenas/3d/motor-three.js +728 -0
- package/editor/modelo-projeto/compositions/cenas/LEIA.md +177 -0
- package/editor/modelo-projeto/ferramentas/medir_fontes.py +65 -0
- package/editor/modelo-projeto/ferramentas/pessoa_recortada.py +176 -0
- package/editor/modelo-projeto/ferramentas/rastrear_maos.py +189 -0
- package/editor/modelo-projeto/hyperframes.json +12 -0
- package/editor/modelo-projeto/package.json +11 -0
- package/editor/modelo-projeto/projeto.json +16 -0
- package/editor/preview/app.css +1009 -0
- package/editor/preview/app.js +2122 -0
- package/editor/preview/index.html +307 -0
- package/editor/pyproject.toml +23 -0
- package/editor/references/3d.md +236 -0
- package/editor/references/animacao-catalogo.md +672 -0
- package/editor/references/animacao.md +498 -0
- package/editor/references/comandos.md +92 -0
- package/editor/references/corte-seco-como-funciona.md +50 -0
- package/editor/references/corte-seco.md +88 -0
- package/editor/references/corte.md +290 -0
- package/editor/references/direcao-por-blocos.md +218 -0
- package/editor/references/enquadramento.md +654 -0
- package/editor/references/estilos-prontos.md +355 -0
- package/editor/references/estudio.md +100 -0
- package/editor/references/exemplo-eventos.json +67 -0
- package/editor/references/lettering.md +252 -0
- package/editor/references/modelos.md +494 -0
- package/editor/references/phase2.md +1007 -0
- package/editor/references/referencia-aplicar.md +203 -0
- package/editor/references/referencia-medir.md +375 -0
- package/editor/references/revisao.md +335 -0
- package/editor/references/som-e-trilha.md +146 -0
- package/editor/references/tabela-de-sons.md +130 -0
- package/editor/references/variacoes.md +234 -0
- package/editor/uv.lock +1141 -0
- package/lib/act.js +36 -8
- package/lib/catalogo.js +4 -0
- package/lib/chamada-torta.js +81 -0
- package/lib/editor-pro.js +786 -0
- package/lib/entrega.js +58 -0
- package/lib/estudio/conferencia-visual.js +76 -0
- package/lib/estudio/diretor.js +6 -1
- package/lib/estudio/exportar.js +3 -1
- package/lib/estudio/gosto.js +45 -4
- package/lib/estudio/index.js +8 -5
- package/lib/estudio/kit-ae/base.jsx +119 -13
- package/lib/estudio/kit-ae/cenas.jsx +47 -15
- package/lib/estudio/kit-ae/guarda.jsx +274 -0
- package/lib/estudio/kit-ae/primitivas.jsx +180 -24
- package/lib/estudio/maquina-diretor.js +61 -7
- package/lib/estudio/maquina-foto.js +177 -0
- package/lib/estudio/maquina-guarda.js +207 -0
- package/lib/estudio/maquina-layout.js +22 -1
- package/lib/estudio/maquina.js +60 -11
- package/lib/estudio/referencias.js +1 -1
- package/lib/estudio/visao/visao.swift +97 -0
- package/lib/estudio/visao.js +60 -0
- package/lib/local/cerebro.json +192 -6
- package/lib/local/motor.js +5 -2
- package/lib/nuvem.js +11 -1
- package/lib/passagem.js +50 -0
- package/lib/perigo.js +19 -11
- package/lib/protecao.js +134 -0
- package/lib/titulo-conversa.js +46 -0
- package/lib/tools.js +3 -0
- package/lib/ui.js +1 -1
- package/package.json +3 -2
- package/skills/editar-meu-video/SKILL.md +6 -2
- package/skills/editar-video/SKILL.md +113 -0
- package/studio/pessoa.py +52 -0
- package/voz/primo_voz/app.py +43 -4
- package/voz/primo_voz/chat.html +17 -6
- package/voz/primo_voz/painel_dados.py +134 -0
|
@@ -0,0 +1,1342 @@
|
|
|
1
|
+
"""Person matte for the "tela verde" edit: the speaker cut out of every frame, cached per cut range.
|
|
2
|
+
|
|
3
|
+
Engine (chosen by measurement on 16/09/2026 against MODNet and Apple Vision, SPEC Contract 13): Google's
|
|
4
|
+
selfie_multiclass_256x256 (Apache-2.0, model card: https://storage.googleapis.com/mediapipe-assets/Model%20Card%20Multiclass%20Segmentation.pdf)
|
|
5
|
+
converted from TFLite to ONNX (assets/models/selfie_multiclass_256x256.onnx, weights unchanged), run by
|
|
6
|
+
onnxruntime on every system: CPU, or CoreML on a Mac with an Apple chip (25% faster, same matte). person =
|
|
7
|
+
1 - background class (softmax of the 6 logits).
|
|
8
|
+
|
|
9
|
+
Per frame, on the output canvas (1080x1920 for a vertical cut):
|
|
10
|
+
1. inference on square tiles along the long side (2 overlapping 1080x1080 tiles on 9:16), each on its own
|
|
11
|
+
session in parallel, blended in the overlap; levels 0.18..0.75 remove the background haze
|
|
12
|
+
2. half size ("work"): DIS optical-flow guided EMA against the previous matte, trusted by the photometric
|
|
13
|
+
warp error (a fast hand follows the current frame; where the picture did not change at all, a jump of the
|
|
14
|
+
model is flicker and the matte only eases toward it); it starts again at every cut. The first frames of a
|
|
15
|
+
range have no past to lean on, so once frame 12 (SETTINGS["start"]) is done, its matte is carried back
|
|
16
|
+
through them (same flow): where the picture barely changed and that later matte clearly says background,
|
|
17
|
+
a patch of the old room the model still took for the person is taken out (only patches that touch the
|
|
18
|
+
background, never a hole inside the body)
|
|
19
|
+
3. cleanup: the largest body plus the pieces near it (hand, microphone), small holes filled, far blobs out
|
|
20
|
+
4. full size: fast colour guided filter only inside an edge band, then a soft contrast curve
|
|
21
|
+
Cache: <edit>/renders/matte/m_<key>.mkv (FFV1 gray, lossless, full output size) + m_<key>.json per cut range.
|
|
22
|
+
The key has the source signature, the picture range, the decode chain, ENGINE_VERSION and SETTINGS, so a new
|
|
23
|
+
cut reuses every range it kept, and the preview and the final share the files.
|
|
24
|
+
Decode: a range is read from its source video with render_cut's own chain (the frames of cut.mp4 before they
|
|
25
|
+
were encoded). The first and last frames are compared with cut.mp4; when they do not line up (a cut made by an
|
|
26
|
+
older version, a source that changed) or the source is gone, the range is read from cut.mp4 itself.
|
|
27
|
+
|
|
28
|
+
Usage:
|
|
29
|
+
uv run python helpers/matte.py --edit-dir DIR [--video cut.mp4] [--force] [--json]
|
|
30
|
+
uv run python helpers/matte.py --check
|
|
31
|
+
uv run python helpers/matte.py --image FOTO.jpg -o recorte.png
|
|
32
|
+
"""
|
|
33
|
+
|
|
34
|
+
from __future__ import annotations
|
|
35
|
+
|
|
36
|
+
import argparse
|
|
37
|
+
import hashlib
|
|
38
|
+
import json
|
|
39
|
+
import math
|
|
40
|
+
import os
|
|
41
|
+
import subprocess
|
|
42
|
+
import sys
|
|
43
|
+
import threading
|
|
44
|
+
import time
|
|
45
|
+
from collections import deque
|
|
46
|
+
from concurrent.futures import ThreadPoolExecutor
|
|
47
|
+
from dataclasses import dataclass, field
|
|
48
|
+
from pathlib import Path
|
|
49
|
+
from typing import Callable, Iterator
|
|
50
|
+
|
|
51
|
+
import numpy as np
|
|
52
|
+
|
|
53
|
+
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
|
54
|
+
from common import ASSETS_DIR, file_signature, fmt_time, load_json, now_iso, probe_video, require_binary, save_json # noqa: E402
|
|
55
|
+
|
|
56
|
+
ENGINE_VERSION = "mc256-3" # bump when the model or the post-processing changes (every cached matte is redone)
|
|
57
|
+
MODEL_NAME = "selfie_multiclass_256x256.onnx"
|
|
58
|
+
MODEL_PATH = ASSETS_DIR / "models" / MODEL_NAME
|
|
59
|
+
MODEL_SHA256 = "35549e14dc62c73b34f0fd04ec1765d9f1577230a2d068cb25c39f600350f477"
|
|
60
|
+
MODEL_SIDE = 256
|
|
61
|
+
SECONDS_PER_SECOND = 3.0 # what to tell the person: about 3 minutes of work per minute of video
|
|
62
|
+
SETTINGS = {
|
|
63
|
+
"tiles": "long_side",
|
|
64
|
+
"levels": [0.18, 0.75],
|
|
65
|
+
"temporal": {"strength": 0.85, "flow_scale": 0.5, "err_sigma": 0.06, "dual": True, "still": [0.012, 0.03]},
|
|
66
|
+
"cleanup": {"near_px": 28, "min_piece": 0.0006, "hole_small": 0.0005, "hole_max": 0.012, "hole_prob": 0.25},
|
|
67
|
+
"refine": {"r": 3, "eps": 0.002, "band_px": 10, "curve": [0.12, 0.88], "erode": 0.6, "gf_div": 3},
|
|
68
|
+
"start": {"frames": 12, "min_range": 8, "err_floor": 0.01, "err_sigma": 0.02, "bg": [0.1, 0.3], "open_px": 6,
|
|
69
|
+
"min_drop": 0.2},
|
|
70
|
+
}
|
|
71
|
+
REF_WORK_SHORT = 540 # the settings above are in px of a 540 wide work frame
|
|
72
|
+
SYNC_FRAMES = 10 # frames compared with cut.mp4 at each end of a range read from the source
|
|
73
|
+
SYNC_SIZE = 96 # short side of the frames compared
|
|
74
|
+
SYNC_MAX_MAD = 2.0 # mean gray difference (0-255) above which the frames are not the same picture
|
|
75
|
+
# (same frame measured at 0.24-0.28 on a real cut; x264 noise is averaged away)
|
|
76
|
+
PROGRESS_EVERY = 5.0 # seconds between progress lines
|
|
77
|
+
CACHE_DIR = Path("renders") / "matte"
|
|
78
|
+
COREML_ENV = "EDITOR_VIDEO_MATTE_COREML"
|
|
79
|
+
THREADS_ENV = "EDITOR_VIDEO_MATTE_THREADS"
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
class MatteError(Exception):
|
|
83
|
+
def __init__(self, message: str, code: int = 1):
|
|
84
|
+
super().__init__(message)
|
|
85
|
+
self.code = code
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
def _cv2():
|
|
89
|
+
os.environ.setdefault("OPENCV_LOG_LEVEL", "ERROR")
|
|
90
|
+
try:
|
|
91
|
+
import cv2 # noqa: PLC0415
|
|
92
|
+
except ImportError as e: # pragma: no cover - depends on the install
|
|
93
|
+
raise MatteError(f"o OpenCV não está instalado ({e}). Rode o instalador de novo (ou 'uv sync' na pasta da skill)", 2)
|
|
94
|
+
return cv2
|
|
95
|
+
|
|
96
|
+
|
|
97
|
+
def _short(obj) -> str:
|
|
98
|
+
return hashlib.sha1(json.dumps(obj, sort_keys=True, ensure_ascii=False).encode("utf-8")).hexdigest()[:16]
|
|
99
|
+
|
|
100
|
+
|
|
101
|
+
# ================================================================== engine check
|
|
102
|
+
|
|
103
|
+
_MODEL_OK: dict[str, str] = {}
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
def import_ort():
|
|
107
|
+
try:
|
|
108
|
+
import onnxruntime as ort # noqa: PLC0415
|
|
109
|
+
except ImportError:
|
|
110
|
+
raise MatteError("a tela verde precisa da biblioteca onnxruntime, que não está instalada. Rode o instalador "
|
|
111
|
+
"de novo (ou 'uv sync' na pasta da skill)", 2) from None
|
|
112
|
+
try:
|
|
113
|
+
ort.set_default_logger_severity(3) # no engine notes on the person's terminal
|
|
114
|
+
except Exception: # noqa: BLE001
|
|
115
|
+
pass
|
|
116
|
+
return ort
|
|
117
|
+
|
|
118
|
+
|
|
119
|
+
def check_model(path: Path = MODEL_PATH) -> None:
|
|
120
|
+
"""The model file is there and is the one this version was measured with."""
|
|
121
|
+
if not path.is_file():
|
|
122
|
+
raise MatteError(f"o modelo do recorte da pessoa não foi encontrado ({path}). Instale a skill de novo", 2)
|
|
123
|
+
sig = file_signature(path)
|
|
124
|
+
if _MODEL_OK.get(str(path)) == sig:
|
|
125
|
+
return
|
|
126
|
+
h = hashlib.sha256()
|
|
127
|
+
with open(path, "rb") as f:
|
|
128
|
+
for chunk in iter(lambda: f.read(1 << 20), b""):
|
|
129
|
+
h.update(chunk)
|
|
130
|
+
if h.hexdigest() != MODEL_SHA256:
|
|
131
|
+
raise MatteError(f"o arquivo do modelo do recorte ({path.name}) está diferente do original (pode ter "
|
|
132
|
+
f"corrompido na cópia). Instale a skill de novo", 2)
|
|
133
|
+
_MODEL_OK[str(path)] = sig
|
|
134
|
+
|
|
135
|
+
|
|
136
|
+
def check_engine() -> dict:
|
|
137
|
+
"""PT-BR errors when the default engine cannot run here. -> {"onnxruntime": version, "model": name}."""
|
|
138
|
+
ort = import_ort()
|
|
139
|
+
check_model()
|
|
140
|
+
_cv2()
|
|
141
|
+
return {"onnxruntime": getattr(ort, "__version__", "?"), "model": MODEL_NAME}
|
|
142
|
+
|
|
143
|
+
|
|
144
|
+
def use_coreml() -> bool:
|
|
145
|
+
"""CoreML on Macs with an Apple chip (measured 16/09/2026: 25% faster, same matte, IoU 0.9997 against the CPU);
|
|
146
|
+
EDITOR_VIDEO_MATTE_COREML=0 keeps the CPU, =1 also tries it on an Intel Mac."""
|
|
147
|
+
raw = (os.environ.get(COREML_ENV) or "").strip()
|
|
148
|
+
if raw in ("0", "1"):
|
|
149
|
+
return raw == "1" and sys.platform == "darwin"
|
|
150
|
+
import platform # noqa: PLC0415
|
|
151
|
+
return sys.platform == "darwin" and platform.machine().lower() in ("arm64", "aarch64")
|
|
152
|
+
|
|
153
|
+
|
|
154
|
+
def _threads_per_session() -> int:
|
|
155
|
+
raw = (os.environ.get(THREADS_ENV) or "").strip()
|
|
156
|
+
if raw.isdigit() and int(raw) > 0:
|
|
157
|
+
return int(raw)
|
|
158
|
+
return max(1, min(4, (os.cpu_count() or 4) // 3))
|
|
159
|
+
|
|
160
|
+
|
|
161
|
+
# ================================================================== model + post-processing
|
|
162
|
+
|
|
163
|
+
def _levels(p: np.ndarray, lo: float, hi: float) -> np.ndarray:
|
|
164
|
+
return np.clip((p - lo) / (hi - lo), 0.0, 1.0)
|
|
165
|
+
|
|
166
|
+
|
|
167
|
+
def smoothstep(x: np.ndarray, e0: float, e1: float) -> np.ndarray:
|
|
168
|
+
t = np.clip((x - e0) / (e1 - e0), 0.0, 1.0)
|
|
169
|
+
return t * t * (3 - 2 * t)
|
|
170
|
+
|
|
171
|
+
|
|
172
|
+
def tile_origins(length: int, side: int) -> list[int]:
|
|
173
|
+
"""Start of each square tile along one axis: as few as cover it, evenly spread."""
|
|
174
|
+
if length <= side:
|
|
175
|
+
return [0]
|
|
176
|
+
n = int(math.ceil(length / side))
|
|
177
|
+
n = max(2, n)
|
|
178
|
+
return [int(round((length - side) * k / (n - 1))) for k in range(n)]
|
|
179
|
+
|
|
180
|
+
|
|
181
|
+
class OnnxSegmenter:
|
|
182
|
+
"""selfie_multiclass_256x256 on onnxruntime: person probability of a frame, tile by tile."""
|
|
183
|
+
|
|
184
|
+
name = "onnx"
|
|
185
|
+
|
|
186
|
+
def __init__(self, sessions: int = 2, threads: int | None = None, coreml: bool | None = None):
|
|
187
|
+
ort = import_ort()
|
|
188
|
+
check_model()
|
|
189
|
+
threads = threads or _threads_per_session()
|
|
190
|
+
if coreml is None:
|
|
191
|
+
coreml = use_coreml()
|
|
192
|
+
cpu: list = ["CPUExecutionProvider"]
|
|
193
|
+
provs = cpu
|
|
194
|
+
if coreml and sys.platform == "darwin" and "CoreMLExecutionProvider" in ort.get_available_providers():
|
|
195
|
+
provs = [("CoreMLExecutionProvider", {"ModelFormat": "NeuralNetwork", "MLComputeUnits": "ALL"})] + cpu
|
|
196
|
+
self.sessions = self._open(ort, sessions, threads, provs)
|
|
197
|
+
if provs is not cpu:
|
|
198
|
+
try: # CoreML compiles on the first run: a failure there falls back to the CPU
|
|
199
|
+
for s in self.sessions:
|
|
200
|
+
s.run(None, {s.get_inputs()[0].name: np.zeros((1, MODEL_SIDE, MODEL_SIDE, 3), np.float32)})
|
|
201
|
+
except Exception: # noqa: BLE001
|
|
202
|
+
self.sessions = self._open(ort, sessions, threads, cpu)
|
|
203
|
+
self.provider = self.sessions[0].get_providers()[0]
|
|
204
|
+
self.input_name = self.sessions[0].get_inputs()[0].name
|
|
205
|
+
self.pool = ThreadPoolExecutor(max_workers=len(self.sessions))
|
|
206
|
+
self.infer_time = 0.0
|
|
207
|
+
self.threads = threads
|
|
208
|
+
|
|
209
|
+
@staticmethod
|
|
210
|
+
def _open(ort, sessions: int, threads: int, provs: list) -> list:
|
|
211
|
+
out = []
|
|
212
|
+
for _ in range(max(1, sessions)):
|
|
213
|
+
so = ort.SessionOptions()
|
|
214
|
+
so.log_severity_level = 3
|
|
215
|
+
so.intra_op_num_threads = threads
|
|
216
|
+
so.inter_op_num_threads = 1
|
|
217
|
+
try:
|
|
218
|
+
out.append(ort.InferenceSession(str(MODEL_PATH), so, providers=provs))
|
|
219
|
+
except Exception as e: # noqa: BLE001
|
|
220
|
+
if provs[0] != "CPUExecutionProvider":
|
|
221
|
+
return OnnxSegmenter._open(ort, sessions, threads, ["CPUExecutionProvider"])
|
|
222
|
+
raise MatteError(f"o onnxruntime não conseguiu abrir o modelo do recorte ({e.__class__.__name__}). "
|
|
223
|
+
f"Instale a skill de novo", 2) from None
|
|
224
|
+
return out
|
|
225
|
+
|
|
226
|
+
def raw(self, tile_rgb: np.ndarray, k: int = 0) -> np.ndarray:
|
|
227
|
+
"""(256, 256, 6) logits of one square RGB tile (any size, resized here)."""
|
|
228
|
+
cv2 = _cv2()
|
|
229
|
+
x = cv2.resize(tile_rgb, (MODEL_SIDE, MODEL_SIDE), interpolation=cv2.INTER_AREA).astype(np.float32)
|
|
230
|
+
x = x / 127.5 - 1.0
|
|
231
|
+
return self.sessions[k % len(self.sessions)].run(None, {self.input_name: x[None]})[0][0]
|
|
232
|
+
|
|
233
|
+
def _person(self, tile_rgb: np.ndarray, k: int, out_size: tuple[int, int]) -> np.ndarray:
|
|
234
|
+
cv2 = _cv2()
|
|
235
|
+
y = self.raw(tile_rgb, k)
|
|
236
|
+
e = np.exp(y - y.max(axis=-1, keepdims=True))
|
|
237
|
+
bg = e[..., 0] / e.sum(axis=-1)
|
|
238
|
+
return cv2.resize((1.0 - bg).astype(np.float32), out_size, interpolation=cv2.INTER_LINEAR)
|
|
239
|
+
|
|
240
|
+
def person_prob(self, rgb: np.ndarray, work: tuple[int, int]) -> np.ndarray:
|
|
241
|
+
"""rgb: HxWx3 uint8 -> float32 person probability at work (w, h)."""
|
|
242
|
+
t0 = time.perf_counter()
|
|
243
|
+
h, w = rgb.shape[:2]
|
|
244
|
+
ww, wh = work
|
|
245
|
+
vertical = h >= w
|
|
246
|
+
side = w if vertical else h
|
|
247
|
+
length, total = (h, wh) if vertical else (w, ww)
|
|
248
|
+
sc = total / length # work px per frame px along the long side
|
|
249
|
+
origins = tile_origins(length, side)
|
|
250
|
+
span = min(total, max(1, int(round(side * sc)))) # tile length at work res
|
|
251
|
+
starts = [min(total - span, int(round(o * sc))) for o in origins]
|
|
252
|
+
jobs = []
|
|
253
|
+
for k, o in enumerate(origins):
|
|
254
|
+
tile = rgb[o:o + side] if vertical else rgb[:, o:o + side]
|
|
255
|
+
size = (ww, span) if vertical else (span, wh)
|
|
256
|
+
jobs.append(self.pool.submit(self._person, tile, k, size))
|
|
257
|
+
acc = np.zeros(total, np.float32)
|
|
258
|
+
out = np.zeros((wh, ww), np.float32)
|
|
259
|
+
for k, (start, job) in enumerate(zip(starts, jobs)):
|
|
260
|
+
p = job.result()
|
|
261
|
+
ramp = np.ones(span, np.float32)
|
|
262
|
+
if k > 0: # overlap with the previous tile: fade in
|
|
263
|
+
ov = max(0, starts[k - 1] + span - start)
|
|
264
|
+
if ov:
|
|
265
|
+
ramp[:ov] = np.linspace(0, 1, ov, dtype=np.float32)
|
|
266
|
+
if k < len(starts) - 1: # overlap with the next tile: fade out
|
|
267
|
+
ov = max(0, start + span - starts[k + 1])
|
|
268
|
+
if ov:
|
|
269
|
+
ramp[span - ov:] = np.minimum(ramp[span - ov:], np.linspace(1, 0, ov, dtype=np.float32))
|
|
270
|
+
acc[start:start + span] += ramp
|
|
271
|
+
if vertical:
|
|
272
|
+
out[start:start + span] += p * ramp[:, None]
|
|
273
|
+
else:
|
|
274
|
+
out[:, start:start + span] += p * ramp[None, :]
|
|
275
|
+
weight = np.maximum(acc, 1e-6)
|
|
276
|
+
out /= weight[:, None] if vertical else weight[None, :]
|
|
277
|
+
self.infer_time += time.perf_counter() - t0
|
|
278
|
+
return out
|
|
279
|
+
|
|
280
|
+
def close(self) -> None:
|
|
281
|
+
self.pool.shutdown(wait=False)
|
|
282
|
+
|
|
283
|
+
|
|
284
|
+
class Temporal:
|
|
285
|
+
"""Optical-flow guided EMA at work resolution."""
|
|
286
|
+
|
|
287
|
+
def __init__(self, strength: float, flow_scale: float, err_sigma: float, dual: bool = True,
|
|
288
|
+
still: tuple[float, float] | None = None):
|
|
289
|
+
cv2 = _cv2()
|
|
290
|
+
self.strength, self.flow_scale, self.err_sigma, self.dual = strength, flow_scale, err_sigma, dual
|
|
291
|
+
self.still = tuple(still) if still else None
|
|
292
|
+
self.dis = cv2.DISOpticalFlow_create(cv2.DISOPTICAL_FLOW_PRESET_FAST)
|
|
293
|
+
self.prev_gray = None
|
|
294
|
+
self.prev_alpha = None
|
|
295
|
+
self._grid = None
|
|
296
|
+
|
|
297
|
+
def reset(self) -> None:
|
|
298
|
+
self.prev_gray = self.prev_alpha = None
|
|
299
|
+
|
|
300
|
+
def carry(self, gray: np.ndarray, prev_gray: np.ndarray, prev_alpha: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
|
|
301
|
+
"""prev_alpha moved onto the frame `gray` (optical flow from prev_gray) -> (alpha, photometric error 0..1)."""
|
|
302
|
+
cv2 = _cv2()
|
|
303
|
+
h, w = gray.shape
|
|
304
|
+
fw, fh = max(8, int(w * self.flow_scale)), max(8, int(h * self.flow_scale))
|
|
305
|
+
g1 = cv2.resize(gray, (fw, fh), interpolation=cv2.INTER_AREA)
|
|
306
|
+
g0 = cv2.resize(prev_gray, (fw, fh), interpolation=cv2.INTER_AREA)
|
|
307
|
+
flow = self.dis.calc(g1, g0, None) # cur(x) ~ prev(x + flow)
|
|
308
|
+
flow = cv2.resize(flow, (w, h), interpolation=cv2.INTER_LINEAR) * (w / fw)
|
|
309
|
+
if self._grid is None or self._grid[0].shape != (h, w):
|
|
310
|
+
self._grid = np.meshgrid(np.arange(w, dtype=np.float32), np.arange(h, dtype=np.float32))
|
|
311
|
+
gx, gy = self._grid
|
|
312
|
+
mx, my = gx + flow[..., 0], gy + flow[..., 1]
|
|
313
|
+
warped_a = cv2.remap(prev_alpha, mx, my, cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
|
|
314
|
+
warped_g = cv2.remap(prev_gray, mx, my, cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)
|
|
315
|
+
g = gray.astype(np.float32)
|
|
316
|
+
err_w = cv2.GaussianBlur(np.abs(g - warped_g.astype(np.float32)) / 255.0, (0, 0), 2)
|
|
317
|
+
if not self.dual:
|
|
318
|
+
return warped_a, err_w
|
|
319
|
+
# where the picture did not move, the previous matte as it was is the better guess
|
|
320
|
+
err_0 = cv2.GaussianBlur(np.abs(g - prev_gray.astype(np.float32)) / 255.0, (0, 0), 2)
|
|
321
|
+
use0 = err_0 <= err_w
|
|
322
|
+
return np.where(use0, prev_alpha, warped_a).astype(np.float32), np.minimum(err_0, err_w)
|
|
323
|
+
|
|
324
|
+
def __call__(self, gray: np.ndarray, p: np.ndarray) -> np.ndarray:
|
|
325
|
+
if self.prev_gray is None or self.strength <= 0:
|
|
326
|
+
self.prev_gray, self.prev_alpha = gray, p
|
|
327
|
+
return p
|
|
328
|
+
prev_a, err = self.carry(gray, self.prev_gray, self.prev_alpha)
|
|
329
|
+
conf = np.exp(-err / self.err_sigma)
|
|
330
|
+
# a big disagreement the flow cannot explain is a real change: follow the model. Where the picture did
|
|
331
|
+
# not change at all, nothing came in or went out: a jump of the model there is its own flicker (a patch
|
|
332
|
+
# of room it takes for the body for a few frames), so the matte keeps easing instead of jumping
|
|
333
|
+
dis = np.abs(p - prev_a)
|
|
334
|
+
jump = smoothstep(dis, 0.35, 0.8)
|
|
335
|
+
if self.still:
|
|
336
|
+
jump = jump * smoothstep(err, self.still[0], self.still[1])
|
|
337
|
+
wprev = self.strength * conf * (1.0 - jump)
|
|
338
|
+
out = ((1 - wprev) * p + wprev * prev_a).astype(np.float32)
|
|
339
|
+
self.prev_gray, self.prev_alpha = gray, out
|
|
340
|
+
return out
|
|
341
|
+
|
|
342
|
+
|
|
343
|
+
def cleanup(a: np.ndarray, prob: np.ndarray | None, near_px: float, min_piece: float, hole_small: float,
|
|
344
|
+
hole_max: float, hole_prob: float) -> np.ndarray:
|
|
345
|
+
"""Keep the body and the pieces close to it; fill the holes the model was unsure about.
|
|
346
|
+
|
|
347
|
+
A background gap the model saw clearly (between an arm and the body, behind glasses) stays open: only
|
|
348
|
+
tiny holes, or holes whose mean model probability is at least hole_prob, are filled."""
|
|
349
|
+
cv2 = _cv2()
|
|
350
|
+
h, w = a.shape
|
|
351
|
+
b = (a > 0.5).astype(np.uint8)
|
|
352
|
+
n, lab, stats, _ = cv2.connectedComponentsWithStats(b, connectivity=8)
|
|
353
|
+
if n <= 1:
|
|
354
|
+
return a
|
|
355
|
+
areas = stats[1:, cv2.CC_STAT_AREA]
|
|
356
|
+
main = 1 + int(np.argmax(areas))
|
|
357
|
+
keep = (lab == main).astype(np.uint8)
|
|
358
|
+
if n > 2:
|
|
359
|
+
dist = cv2.distanceTransform(1 - keep, cv2.DIST_L2, 3)
|
|
360
|
+
for i in range(1, n):
|
|
361
|
+
if i == main or stats[i, cv2.CC_STAT_AREA] < min_piece * h * w:
|
|
362
|
+
continue
|
|
363
|
+
comp = lab == i
|
|
364
|
+
if dist[comp].min() <= near_px:
|
|
365
|
+
keep[comp] = 1
|
|
366
|
+
inv = (1 - keep).astype(np.uint8)
|
|
367
|
+
n2, lab2, st2, _ = cv2.connectedComponentsWithStats(inv, connectivity=4)
|
|
368
|
+
holes = np.zeros_like(keep)
|
|
369
|
+
src = a if prob is None else prob
|
|
370
|
+
for i in range(1, n2):
|
|
371
|
+
x, y, ww, hh, ar = st2[i]
|
|
372
|
+
touches = x == 0 or y == 0 or x + ww == w or y + hh == h
|
|
373
|
+
if touches or ar >= hole_max * h * w:
|
|
374
|
+
continue
|
|
375
|
+
comp = lab2 == i
|
|
376
|
+
if ar < hole_small * h * w or float(src[comp].mean()) >= hole_prob:
|
|
377
|
+
holes[comp] = 1
|
|
378
|
+
solid = keep | holes
|
|
379
|
+
# soft keep mask: 1 on the kept pieces and 4 px around them, fading to 0 by 8 px
|
|
380
|
+
dout = cv2.distanceTransform(1 - solid, cv2.DIST_L2, 3)
|
|
381
|
+
keep_soft = np.clip((8.0 - dout) / 4.0, 0, 1)
|
|
382
|
+
out = a * keep_soft
|
|
383
|
+
if holes.any():
|
|
384
|
+
out = np.maximum(out, cv2.GaussianBlur(holes.astype(np.float32), (0, 0), 1.0) * holes)
|
|
385
|
+
return out.astype(np.float32)
|
|
386
|
+
|
|
387
|
+
|
|
388
|
+
def _box(x, r):
|
|
389
|
+
cv2 = _cv2()
|
|
390
|
+
return cv2.boxFilter(x, -1, (2 * r + 1, 2 * r + 1), borderType=cv2.BORDER_REFLECT)
|
|
391
|
+
|
|
392
|
+
|
|
393
|
+
def fast_guided_filter(guide_full: np.ndarray, p_low: np.ndarray, r: int, eps: float) -> np.ndarray:
|
|
394
|
+
"""Colour guided filter (He et al.): coefficients at the size of p_low, applied with the full size guide."""
|
|
395
|
+
cv2 = _cv2()
|
|
396
|
+
H, W = guide_full.shape[:2]
|
|
397
|
+
h, w = p_low.shape
|
|
398
|
+
I = cv2.resize(guide_full, (w, h), interpolation=cv2.INTER_AREA) # noqa: E741
|
|
399
|
+
p = p_low
|
|
400
|
+
mI = _box(I, r)
|
|
401
|
+
mp_ = _box(p, r)
|
|
402
|
+
mIp = _box(I * p[..., None], r)
|
|
403
|
+
cov = mIp - mI * mp_[..., None]
|
|
404
|
+
rr = _box(I[..., 0] * I[..., 0], r) - mI[..., 0] * mI[..., 0] + eps
|
|
405
|
+
rg = _box(I[..., 0] * I[..., 1], r) - mI[..., 0] * mI[..., 1]
|
|
406
|
+
rb = _box(I[..., 0] * I[..., 2], r) - mI[..., 0] * mI[..., 2]
|
|
407
|
+
gg = _box(I[..., 1] * I[..., 1], r) - mI[..., 1] * mI[..., 1] + eps
|
|
408
|
+
gb = _box(I[..., 1] * I[..., 2], r) - mI[..., 1] * mI[..., 2]
|
|
409
|
+
bb = _box(I[..., 2] * I[..., 2], r) - mI[..., 2] * mI[..., 2] + eps
|
|
410
|
+
inv_rr = gg * bb - gb * gb
|
|
411
|
+
inv_rg = gb * rb - rg * bb
|
|
412
|
+
inv_rb = rg * gb - gg * rb
|
|
413
|
+
inv_gg = rr * bb - rb * rb
|
|
414
|
+
inv_gb = rb * rg - rr * gb
|
|
415
|
+
inv_bb = rr * gg - rg * rg
|
|
416
|
+
det = rr * inv_rr + rg * inv_rg + rb * inv_rb
|
|
417
|
+
c0, c1, c2 = cov[..., 0], cov[..., 1], cov[..., 2]
|
|
418
|
+
a0 = (inv_rr * c0 + inv_rg * c1 + inv_rb * c2) / det
|
|
419
|
+
a1 = (inv_rg * c0 + inv_gg * c1 + inv_gb * c2) / det
|
|
420
|
+
a2 = (inv_rb * c0 + inv_gb * c1 + inv_bb * c2) / det
|
|
421
|
+
A = np.dstack([a0, a1, a2]).astype(np.float32)
|
|
422
|
+
b = (mp_ - (A * mI).sum(axis=2)).astype(np.float32)
|
|
423
|
+
A = _box(A, r)
|
|
424
|
+
b = _box(b, r)
|
|
425
|
+
A_full = cv2.resize(A, (W, H), interpolation=cv2.INTER_LINEAR)
|
|
426
|
+
q = cv2.resize(b, (W, H), interpolation=cv2.INTER_LINEAR)
|
|
427
|
+
for c in range(3):
|
|
428
|
+
q += cv2.multiply(A_full[..., c], guide_full[..., c])
|
|
429
|
+
return q
|
|
430
|
+
|
|
431
|
+
|
|
432
|
+
class Refiner:
|
|
433
|
+
def __init__(self, r: int, eps: float, band_px: float, curve, erode: float, gf_div: int):
|
|
434
|
+
self.r, self.eps, self.band_px, self.curve, self.erode, self.gf_div = r, eps, band_px, tuple(curve), erode, gf_div
|
|
435
|
+
|
|
436
|
+
def __call__(self, rgb_full: np.ndarray, a_low: np.ndarray) -> np.ndarray:
|
|
437
|
+
cv2 = _cv2()
|
|
438
|
+
H, W = rgb_full.shape[:2]
|
|
439
|
+
guide = rgb_full.astype(np.float32) / 255.0
|
|
440
|
+
gf_size = (max(8, W // self.gf_div), max(8, H // self.gf_div))
|
|
441
|
+
p_gf = cv2.resize(a_low, gf_size, interpolation=cv2.INTER_AREA)
|
|
442
|
+
q = fast_guided_filter(guide, p_gf, self.r, self.eps)
|
|
443
|
+
coarse = cv2.resize(a_low, (W, H), interpolation=cv2.INTER_LINEAR)
|
|
444
|
+
b = (a_low > 0.5).astype(np.uint8)
|
|
445
|
+
d_in = cv2.distanceTransform(b, cv2.DIST_L2, 3)
|
|
446
|
+
d_out = cv2.distanceTransform(1 - b, cv2.DIST_L2, 3)
|
|
447
|
+
d = np.maximum(d_in, d_out) # distance to the contour (work px)
|
|
448
|
+
band = np.clip((self.band_px + 2 - d) / 4.0, 0, 1).astype(np.float32)
|
|
449
|
+
band = cv2.resize(band, (W, H), interpolation=cv2.INTER_LINEAR)
|
|
450
|
+
a = band * np.clip(q, 0, 1) + (1 - band) * coarse
|
|
451
|
+
a = smoothstep(a, *self.curve)
|
|
452
|
+
if self.erode > 0: # pull the edge in a little so the old background does not fringe
|
|
453
|
+
a = cv2.erode(a, np.ones((3, 3), np.uint8)) * self.erode + a * (1 - self.erode)
|
|
454
|
+
return a.astype(np.float32)
|
|
455
|
+
|
|
456
|
+
|
|
457
|
+
class MatteEngine:
|
|
458
|
+
"""Frame by frame person matte. process() must see the frames of one range in order; reset() between ranges."""
|
|
459
|
+
|
|
460
|
+
version = ENGINE_VERSION
|
|
461
|
+
|
|
462
|
+
def __init__(self, size: tuple[int, int], segmenter: OnnxSegmenter | None = None):
|
|
463
|
+
W, H = size
|
|
464
|
+
self.size = (W, H)
|
|
465
|
+
self.work = (max(2, (W // 2) // 2 * 2), max(2, (H // 2) // 2 * 2))
|
|
466
|
+
k = min(self.work) / REF_WORK_SHORT
|
|
467
|
+
t, c, r = SETTINGS["temporal"], SETTINGS["cleanup"], SETTINGS["refine"]
|
|
468
|
+
self.seg = segmenter or OnnxSegmenter()
|
|
469
|
+
self.provider = self.seg.provider
|
|
470
|
+
self.temporal = Temporal(t["strength"], t["flow_scale"], t["err_sigma"], t["dual"], t.get("still"))
|
|
471
|
+
self.clean = dict(c, near_px=c["near_px"] * k)
|
|
472
|
+
self.refiner = Refiner(r["r"], r["eps"], r["band_px"] * k, r["curve"], r["erode"], r["gf_div"])
|
|
473
|
+
|
|
474
|
+
def reset(self) -> None:
|
|
475
|
+
self.temporal.reset()
|
|
476
|
+
|
|
477
|
+
def coarse_parts(self, rgb: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
|
478
|
+
"""(temporal matte, model probability, gray) of the next frame at work size (sequential)."""
|
|
479
|
+
cv2 = _cv2()
|
|
480
|
+
p = self.seg.person_prob(rgb, self.work)
|
|
481
|
+
lo, hi = SETTINGS["levels"]
|
|
482
|
+
p = _levels(p, lo, hi)
|
|
483
|
+
gray = cv2.cvtColor(cv2.resize(rgb, self.work, interpolation=cv2.INTER_AREA), cv2.COLOR_RGB2GRAY)
|
|
484
|
+
return self.temporal(gray, p), p, gray
|
|
485
|
+
|
|
486
|
+
def finish(self, a: np.ndarray, p: np.ndarray) -> np.ndarray:
|
|
487
|
+
return cleanup(a, p, **self.clean)
|
|
488
|
+
|
|
489
|
+
def coarse(self, rgb: np.ndarray) -> np.ndarray:
|
|
490
|
+
"""Work-size matte of the next frame (sequential: it feeds the temporal filter)."""
|
|
491
|
+
a, p, _ = self.coarse_parts(rgb)
|
|
492
|
+
return self.finish(a, p)
|
|
493
|
+
|
|
494
|
+
def settle_start(self, first: list[tuple[np.ndarray, np.ndarray, np.ndarray]],
|
|
495
|
+
anchor: tuple[np.ndarray, np.ndarray]) -> list[np.ndarray]:
|
|
496
|
+
"""The first frames of a range, (temporal matte, probability, gray) each, with the matte and gray of the
|
|
497
|
+
frame right after them (anchor): the anchor matte is carried back frame by frame; where the picture
|
|
498
|
+
barely changed along the way and it clearly says background, the start matte drops to it. Only
|
|
499
|
+
patches of START open_px or more that touch the background change (a thin moving hand, a hole inside
|
|
500
|
+
the body never do), fading in over the last half before the anchor. -> corrected temporal mattes."""
|
|
501
|
+
cv2 = _cv2()
|
|
502
|
+
st = SETTINGS["start"]
|
|
503
|
+
n = len(first)
|
|
504
|
+
if n == 0:
|
|
505
|
+
return []
|
|
506
|
+
b, gprev = anchor
|
|
507
|
+
k = min(self.work) / REF_WORK_SHORT
|
|
508
|
+
r = max(1, int(round(st["open_px"] * k)))
|
|
509
|
+
ker = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * r + 1, 2 * r + 1))
|
|
510
|
+
near = np.ones((5, 5), np.uint8)
|
|
511
|
+
cerr = None
|
|
512
|
+
out: list[np.ndarray] = [None] * n # type: ignore[list-item]
|
|
513
|
+
for i in range(n - 1, -1, -1):
|
|
514
|
+
a, _, gray = first[i]
|
|
515
|
+
b, err = self.temporal.carry(gray, gprev, b)
|
|
516
|
+
gprev = gray
|
|
517
|
+
cerr = err if cerr is None else np.maximum(cerr, err)
|
|
518
|
+
w = np.exp(-np.maximum(0.0, cerr - st["err_floor"]) / st["err_sigma"])
|
|
519
|
+
w *= 1.0 - smoothstep(b, *st["bg"])
|
|
520
|
+
w *= min(1.0, (n - i) / (n / 2.0))
|
|
521
|
+
drop = (w * np.maximum(0.0, a - b)).astype(np.float32)
|
|
522
|
+
m = cv2.morphologyEx((drop > st["min_drop"]).astype(np.uint8), cv2.MORPH_OPEN, ker)
|
|
523
|
+
if m.any():
|
|
524
|
+
cnt, lab = cv2.connectedComponents(m, connectivity=8)
|
|
525
|
+
touch = cv2.dilate((a < 0.5).astype(np.uint8), near) > 0
|
|
526
|
+
keep = np.zeros(cnt, bool)
|
|
527
|
+
keep[np.unique(lab[touch & (m > 0)])] = True
|
|
528
|
+
keep[0] = False
|
|
529
|
+
m = keep[lab].astype(np.uint8)
|
|
530
|
+
if not m.any():
|
|
531
|
+
out[i] = a
|
|
532
|
+
continue
|
|
533
|
+
soft = cv2.GaussianBlur(cv2.dilate(m, ker).astype(np.float32), (0, 0), r / 2.0)
|
|
534
|
+
out[i] = (a - drop * soft).astype(np.float32)
|
|
535
|
+
return out
|
|
536
|
+
|
|
537
|
+
def refine(self, rgb: np.ndarray, coarse: np.ndarray) -> np.ndarray:
|
|
538
|
+
return self.refiner(rgb, coarse)
|
|
539
|
+
|
|
540
|
+
def process(self, rgb: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
|
|
541
|
+
a = self.coarse(rgb)
|
|
542
|
+
return a, self.refine(rgb, a)
|
|
543
|
+
|
|
544
|
+
def close(self) -> None:
|
|
545
|
+
self.seg.close()
|
|
546
|
+
|
|
547
|
+
|
|
548
|
+
# Engines can be swapped (tests use a fake one so they do not depend on the model).
|
|
549
|
+
EngineFactory = Callable[[tuple[int, int]], object]
|
|
550
|
+
_ENGINE_FACTORY: EngineFactory | None = None
|
|
551
|
+
_ENGINE_VERSION_OVERRIDE: str | None = None
|
|
552
|
+
|
|
553
|
+
|
|
554
|
+
def set_engine(factory: EngineFactory | None, version: str | None = None) -> None:
|
|
555
|
+
"""Use another engine (an object with version, provider, reset(), coarse(rgb), refine(rgb, coarse), close()).
|
|
556
|
+
None goes back to the ONNX engine."""
|
|
557
|
+
global _ENGINE_FACTORY, _ENGINE_VERSION_OVERRIDE
|
|
558
|
+
_ENGINE_FACTORY = factory
|
|
559
|
+
_ENGINE_VERSION_OVERRIDE = version if factory is not None else None
|
|
560
|
+
|
|
561
|
+
|
|
562
|
+
def engine_version() -> str:
|
|
563
|
+
return _ENGINE_VERSION_OVERRIDE or ENGINE_VERSION
|
|
564
|
+
|
|
565
|
+
|
|
566
|
+
def make_engine(size: tuple[int, int]):
|
|
567
|
+
if _ENGINE_FACTORY is not None:
|
|
568
|
+
return _ENGINE_FACTORY(size)
|
|
569
|
+
return MatteEngine(size)
|
|
570
|
+
|
|
571
|
+
|
|
572
|
+
def default_engine_ready() -> None:
|
|
573
|
+
"""Raise MatteError when the engine that would run cannot (a fake engine is always ready)."""
|
|
574
|
+
if _ENGINE_FACTORY is None:
|
|
575
|
+
check_engine()
|
|
576
|
+
|
|
577
|
+
|
|
578
|
+
# ================================================================== stats
|
|
579
|
+
|
|
580
|
+
def frame_stats(a: np.ndarray) -> tuple | None:
|
|
581
|
+
"""Per frame numbers used to place the person (work-size matte): None when nobody is there.
|
|
582
|
+
|
|
583
|
+
-> (top, left, right, head_cx, touch_left, touch_right, touch_bottom) as fractions of the frame / bools."""
|
|
584
|
+
m = a > 0.5
|
|
585
|
+
rows = np.flatnonzero(m.any(axis=1))
|
|
586
|
+
if rows.size == 0:
|
|
587
|
+
return None
|
|
588
|
+
h, w = a.shape
|
|
589
|
+
cols = np.flatnonzero(m.any(axis=0))
|
|
590
|
+
r0 = int(rows[0])
|
|
591
|
+
r1 = r0 + max(1, int((h - r0) * 0.35))
|
|
592
|
+
band = a[r0:r1]
|
|
593
|
+
mass = band.sum(axis=0)
|
|
594
|
+
cx = float((mass * np.arange(w)).sum() / max(1e-6, mass.sum()))
|
|
595
|
+
edge = max(1, w // 270)
|
|
596
|
+
return (r0 / h, int(cols[0]) / w, (int(cols[-1]) + 1) / w, cx / w,
|
|
597
|
+
bool(m[:, :edge].any()), bool(m[:, w - edge:].any()), bool(m[h - max(1, h // 480):].any()))
|
|
598
|
+
|
|
599
|
+
|
|
600
|
+
def summarize_stats(rows: list[tuple | None]) -> dict:
|
|
601
|
+
found = [r for r in rows if r is not None]
|
|
602
|
+
n = len(rows)
|
|
603
|
+
if not found:
|
|
604
|
+
return {"frames": n, "person_frames": 0}
|
|
605
|
+
arr = np.array([r[:4] for r in found], np.float64)
|
|
606
|
+
touch = np.array([r[4:] for r in found], np.float64)
|
|
607
|
+
return {
|
|
608
|
+
"frames": n,
|
|
609
|
+
"person_frames": len(found),
|
|
610
|
+
"top": round(float(np.percentile(arr[:, 0], 5)), 5),
|
|
611
|
+
"top_median": round(float(np.median(arr[:, 0])), 5),
|
|
612
|
+
"left": round(float(np.percentile(arr[:, 1], 5)), 5),
|
|
613
|
+
"right": round(float(np.percentile(arr[:, 2], 95)), 5),
|
|
614
|
+
"head_cx": round(float(np.median(arr[:, 3])), 5),
|
|
615
|
+
"touch_left": round(float(touch[:, 0].mean()), 4),
|
|
616
|
+
"touch_right": round(float(touch[:, 1].mean()), 4),
|
|
617
|
+
"touch_bottom": round(float(touch[:, 2].mean()), 4),
|
|
618
|
+
}
|
|
619
|
+
|
|
620
|
+
|
|
621
|
+
# ================================================================== ranges
|
|
622
|
+
|
|
623
|
+
@dataclass
|
|
624
|
+
class RangeJob:
|
|
625
|
+
index: int
|
|
626
|
+
rid: str
|
|
627
|
+
f0: int # first frame on the base video (output frame grid)
|
|
628
|
+
f1: int
|
|
629
|
+
size: tuple[int, int] # full output size (the matte is stored at this size)
|
|
630
|
+
fps: int
|
|
631
|
+
key: str = ""
|
|
632
|
+
method: str = "corte" # "fonte" (source video) | "corte" (cut.mp4)
|
|
633
|
+
source: Path | None = None
|
|
634
|
+
source_args: list[str] = field(default_factory=list) # ffmpeg input + filter args (method fonte)
|
|
635
|
+
source_hw: bool = False
|
|
636
|
+
base: Path | None = None
|
|
637
|
+
base_pre: list[str] = field(default_factory=list) # filters that put the base on the output grid
|
|
638
|
+
base_fps: float = 0.0
|
|
639
|
+
note: str = "" # why the range is read from cut.mp4
|
|
640
|
+
path: Path = field(default_factory=Path)
|
|
641
|
+
meta_path: Path = field(default_factory=Path)
|
|
642
|
+
meta: dict = field(default_factory=dict)
|
|
643
|
+
cached: bool = False
|
|
644
|
+
seconds: float = 0.0
|
|
645
|
+
|
|
646
|
+
@property
|
|
647
|
+
def frames(self) -> int:
|
|
648
|
+
return self.f1 - self.f0
|
|
649
|
+
|
|
650
|
+
|
|
651
|
+
def base_pre_filters(base_fps: float, base_size: tuple[int, int], fps: int, size: tuple[int, int]) -> list[str]:
|
|
652
|
+
"""The same frame-grid filters compose_final.run_video puts on the base."""
|
|
653
|
+
pre = []
|
|
654
|
+
if abs(base_fps - fps) > 0.01:
|
|
655
|
+
pre.append(f"fps={fps}")
|
|
656
|
+
if tuple(base_size) != tuple(size):
|
|
657
|
+
W, H = size
|
|
658
|
+
pre.append(f"scale={W}:{H}:force_original_aspect_ratio=increase:flags=bicubic,crop={W}:{H},setsar=1")
|
|
659
|
+
return pre
|
|
660
|
+
|
|
661
|
+
|
|
662
|
+
def _timeline_segment(timeline: dict | None, f0: int, f1: int, fps: int) -> dict | None:
|
|
663
|
+
for s in (timeline or {}).get("segments") or []:
|
|
664
|
+
try:
|
|
665
|
+
a = int(s["out_start_frame"]) if "out_start_frame" in s else int(round(float(s["out_start"]) * fps))
|
|
666
|
+
b = int(s["out_end_frame"]) if "out_end_frame" in s else int(round(float(s["out_end"]) * fps))
|
|
667
|
+
except (KeyError, TypeError, ValueError):
|
|
668
|
+
continue
|
|
669
|
+
if a == f0 and b == f1:
|
|
670
|
+
return s
|
|
671
|
+
return None
|
|
672
|
+
|
|
673
|
+
|
|
674
|
+
def _source_plan(job: RangeJob, edit_dir: Path, timeline: dict, seg: dict, sources: dict | None,
|
|
675
|
+
base_info: dict) -> str:
|
|
676
|
+
"""Fill the source decode of a job. Returns "" or why the source cannot be used."""
|
|
677
|
+
import render_cut as rc # noqa: PLC0415
|
|
678
|
+
from grade import resolve_grade # noqa: PLC0415
|
|
679
|
+
|
|
680
|
+
W, H = job.size
|
|
681
|
+
try:
|
|
682
|
+
tl_fps = int(round(float(timeline.get("fps"))))
|
|
683
|
+
tl_size = (int(timeline.get("width")), int(timeline.get("height")))
|
|
684
|
+
except (TypeError, ValueError):
|
|
685
|
+
return "o cut_timeline.json não tem fps ou tamanho"
|
|
686
|
+
if timeline.get("preview"):
|
|
687
|
+
return "o corte é uma prévia"
|
|
688
|
+
if tl_fps != job.fps or tl_size != (W, H):
|
|
689
|
+
return "o corte tem outro tamanho ou outra taxa de quadros"
|
|
690
|
+
if abs(float(base_info.get("fps") or 0) - job.fps) > 0.01 or tuple(base_info.get("size") or ()) != (W, H):
|
|
691
|
+
return "o vídeo base não está no tamanho do corte"
|
|
692
|
+
other = bool(seg.get("imagem_de")) # the picture came from another camera (render_cut "imagem_de")
|
|
693
|
+
try:
|
|
694
|
+
src_start = float(seg["imagem_src_start"] if other else seg["src_start"])
|
|
695
|
+
frames = int(seg.get("frames") or job.frames)
|
|
696
|
+
except (KeyError, TypeError, ValueError):
|
|
697
|
+
return "o cut_timeline.json é de uma versão antiga"
|
|
698
|
+
if frames != job.frames:
|
|
699
|
+
return "o trecho do timeline não bate com o vídeo"
|
|
700
|
+
if other:
|
|
701
|
+
src = Path(str(seg.get("imagem_path") or ""))
|
|
702
|
+
if not src.is_file():
|
|
703
|
+
return "o vídeo bruto da outra câmera não foi encontrado"
|
|
704
|
+
elif not sources or str(seg.get("source")) not in sources:
|
|
705
|
+
return "o vídeo bruto não foi encontrado"
|
|
706
|
+
else:
|
|
707
|
+
src = sources[str(seg.get("source"))]
|
|
708
|
+
try:
|
|
709
|
+
info = probe_video(src)
|
|
710
|
+
except Exception: # noqa: BLE001
|
|
711
|
+
return f"não consegui ler {src.name}"
|
|
712
|
+
if not info.get("has_video"):
|
|
713
|
+
return f"{src.name} não tem imagem"
|
|
714
|
+
try:
|
|
715
|
+
grade_filter = resolve_grade(timeline.get("grade"))
|
|
716
|
+
except ValueError:
|
|
717
|
+
return "o grade do corte não existe mais"
|
|
718
|
+
try:
|
|
719
|
+
vf = rc.build_vf(info, W, H, job.fps, grade_filter, seg.get("crop"))
|
|
720
|
+
except Exception: # noqa: BLE001
|
|
721
|
+
return "não consegui montar a leitura do vídeo bruto"
|
|
722
|
+
snapped = bool(seg.get("pic_snapped", True))
|
|
723
|
+
src_fps = float(info.get("fps") or job.fps)
|
|
724
|
+
seek = max(0.0, src_start - (0.25 / max(src_fps, job.fps) if snapped else 0.0))
|
|
725
|
+
job.method = "fonte"
|
|
726
|
+
job.source = src
|
|
727
|
+
job.source_args = ["-noaccurate_seek", "-ss", f"{seek:.6f}", "-i", str(src), "-map", "0:v:0", "-vf", vf]
|
|
728
|
+
try:
|
|
729
|
+
job.source_hw = bool(rc.hwaccel_for(src, info))
|
|
730
|
+
except Exception: # noqa: BLE001
|
|
731
|
+
job.source_hw = False
|
|
732
|
+
job.key = _short({"v": engine_version(), "s": SETTINGS, "how": "fonte", "src": str(src.resolve()),
|
|
733
|
+
"sig": file_signature(src), "pic": round(src_start, 6), "snap": snapped, "frames": job.frames,
|
|
734
|
+
"vf": vf, "size": [W, H], "fps": job.fps})
|
|
735
|
+
return ""
|
|
736
|
+
|
|
737
|
+
|
|
738
|
+
def plan_jobs(edit_dir: Path, base: Path, bounds: list[tuple[int, int]], fps: int, size: tuple[int, int],
|
|
739
|
+
timeline: dict | None, base_fps: float | None = None, base_size: tuple[int, int] | None = None) -> list[RangeJob]:
|
|
740
|
+
"""One job per cut range (bounds on the output frame grid)."""
|
|
741
|
+
edit_dir = Path(edit_dir).resolve()
|
|
742
|
+
base = Path(base).resolve()
|
|
743
|
+
if base_fps is None or base_size is None:
|
|
744
|
+
info = probe_video(base)
|
|
745
|
+
base_fps = float(info.get("fps") or fps)
|
|
746
|
+
base_size = (int(info.get("display_width") or info.get("width") or 0),
|
|
747
|
+
int(info.get("display_height") or info.get("height") or 0))
|
|
748
|
+
base_info = {"fps": base_fps, "size": tuple(base_size)}
|
|
749
|
+
pre = base_pre_filters(base_fps, base_size, fps, size)
|
|
750
|
+
sources = None
|
|
751
|
+
sources_note = ""
|
|
752
|
+
if timeline and timeline.get("segments"):
|
|
753
|
+
try:
|
|
754
|
+
import denoise # noqa: PLC0415
|
|
755
|
+
sources, sources_note = denoise.resolve_segment_sources(edit_dir, timeline)
|
|
756
|
+
except Exception as e: # noqa: BLE001
|
|
757
|
+
sources, sources_note = None, f"não consegui achar os vídeos brutos ({e.__class__.__name__})"
|
|
758
|
+
base_sig = file_signature(base)
|
|
759
|
+
out_dir = edit_dir / CACHE_DIR
|
|
760
|
+
jobs = []
|
|
761
|
+
for i, (a, b) in enumerate(bounds):
|
|
762
|
+
seg = _timeline_segment(timeline, a, b, fps)
|
|
763
|
+
rid = str(seg.get("id")) if seg and seg.get("id") else f"t{i + 1}"
|
|
764
|
+
job = RangeJob(index=i, rid=rid, f0=a, f1=b, size=tuple(size), fps=fps, base=base, base_pre=pre,
|
|
765
|
+
base_fps=float(base_fps))
|
|
766
|
+
why = "sem cut_timeline.json" if not timeline else ("o trecho não está no cut_timeline.json" if seg is None else "")
|
|
767
|
+
if seg is not None:
|
|
768
|
+
why = sources_note if sources is None and sources_note else ""
|
|
769
|
+
if not why:
|
|
770
|
+
why = _source_plan(job, edit_dir, timeline, seg, sources, base_info)
|
|
771
|
+
if job.method != "fonte":
|
|
772
|
+
job.note = why
|
|
773
|
+
job.key = _short({"v": engine_version(), "s": SETTINGS, "how": "corte", "base": base.name,
|
|
774
|
+
"sig": base_sig, "f0": a, "f1": b, "size": list(size), "fps": fps, "pre": pre})
|
|
775
|
+
job.path = out_dir / f"m_{job.key}.mkv"
|
|
776
|
+
job.meta_path = out_dir / f"m_{job.key}.json"
|
|
777
|
+
jobs.append(job)
|
|
778
|
+
return jobs
|
|
779
|
+
|
|
780
|
+
|
|
781
|
+
# ================================================================== decode / encode
|
|
782
|
+
|
|
783
|
+
class FrameReader:
|
|
784
|
+
"""Raw frames from an ffmpeg command, read in a thread (the decode runs while the engine works)."""
|
|
785
|
+
|
|
786
|
+
def __init__(self, cmd: list[str], size: tuple[int, int], channels: int, what: str, queue_len: int = 6):
|
|
787
|
+
self.size, self.channels, self.what = size, channels, what
|
|
788
|
+
self.proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.DEVNULL)
|
|
789
|
+
self.frame_bytes = size[0] * size[1] * channels
|
|
790
|
+
self.buf: deque = deque()
|
|
791
|
+
self.cond = threading.Condition()
|
|
792
|
+
self.done = False
|
|
793
|
+
self.max = queue_len
|
|
794
|
+
self.err = b""
|
|
795
|
+
self.thread = threading.Thread(target=self._run, daemon=True)
|
|
796
|
+
self.thread.start()
|
|
797
|
+
self.err_thread = threading.Thread(target=self._drain_err, daemon=True)
|
|
798
|
+
self.err_thread.start()
|
|
799
|
+
|
|
800
|
+
def _drain_err(self) -> None:
|
|
801
|
+
try:
|
|
802
|
+
self.err = self.proc.stderr.read() if self.proc.stderr else b""
|
|
803
|
+
except (OSError, ValueError):
|
|
804
|
+
pass
|
|
805
|
+
|
|
806
|
+
def _run(self) -> None:
|
|
807
|
+
W, H = self.size
|
|
808
|
+
shape = (H, W, self.channels) if self.channels > 1 else (H, W)
|
|
809
|
+
try:
|
|
810
|
+
while True:
|
|
811
|
+
data = self.proc.stdout.read(self.frame_bytes) if self.proc.stdout else b""
|
|
812
|
+
if len(data) < self.frame_bytes:
|
|
813
|
+
break
|
|
814
|
+
frame = np.frombuffer(data, np.uint8).reshape(shape)
|
|
815
|
+
with self.cond:
|
|
816
|
+
while len(self.buf) >= self.max and not self.done:
|
|
817
|
+
self.cond.wait()
|
|
818
|
+
if self.done:
|
|
819
|
+
break
|
|
820
|
+
self.buf.append(frame)
|
|
821
|
+
self.cond.notify_all()
|
|
822
|
+
finally:
|
|
823
|
+
with self.cond:
|
|
824
|
+
self.done = True
|
|
825
|
+
self.cond.notify_all()
|
|
826
|
+
|
|
827
|
+
def read(self) -> np.ndarray | None:
|
|
828
|
+
with self.cond:
|
|
829
|
+
while not self.buf and not self.done:
|
|
830
|
+
self.cond.wait()
|
|
831
|
+
if self.buf:
|
|
832
|
+
frame = self.buf.popleft()
|
|
833
|
+
self.cond.notify_all()
|
|
834
|
+
return frame
|
|
835
|
+
return None
|
|
836
|
+
|
|
837
|
+
def close(self) -> str:
|
|
838
|
+
with self.cond:
|
|
839
|
+
self.done = True
|
|
840
|
+
self.buf.clear()
|
|
841
|
+
self.cond.notify_all()
|
|
842
|
+
if self.proc.poll() is None:
|
|
843
|
+
self.proc.kill()
|
|
844
|
+
self.proc.wait()
|
|
845
|
+
self.err_thread.join(timeout=2)
|
|
846
|
+
return self.err.decode("utf-8", "replace").strip()
|
|
847
|
+
|
|
848
|
+
|
|
849
|
+
def _ff() -> list[str]:
|
|
850
|
+
return ["ffmpeg", "-hide_banner", "-nostdin", "-v", "error"]
|
|
851
|
+
|
|
852
|
+
|
|
853
|
+
def source_cmd(job: RangeJob, hw: bool, size: tuple[int, int] | None = None, frames: int | None = None) -> list[str]:
|
|
854
|
+
args = list(job.source_args)
|
|
855
|
+
vf_i = args.index("-vf") + 1
|
|
856
|
+
W, H = size or job.size
|
|
857
|
+
extra = "" if (W, H) == tuple(job.size) else f",scale={W}:{H}:flags=area"
|
|
858
|
+
args[vf_i] = args[vf_i] + extra + ",format=rgb24"
|
|
859
|
+
cmd = _ff() + (["-hwaccel", "videotoolbox"] if hw else []) + args
|
|
860
|
+
return cmd + ["-frames:v", str(frames or job.frames), "-an", "-f", "rawvideo", "-pix_fmt", "rgb24", "-"]
|
|
861
|
+
|
|
862
|
+
|
|
863
|
+
def base_cmd(base: Path, pre: list[str], base_fps: float, fps: int, f0: int, n: int,
|
|
864
|
+
size: tuple[int, int] | None = None, full_size: tuple[int, int] | None = None, pix_fmt: str = "rgb24") -> list[str]:
|
|
865
|
+
"""Frames f0..f0+n of the base on the output grid (exact: accurate seek without an fps change, else from 0)."""
|
|
866
|
+
filters = list(pre)
|
|
867
|
+
if size is not None and full_size is not None and tuple(size) != tuple(full_size):
|
|
868
|
+
filters.append(f"scale={size[0]}:{size[1]}:flags=area")
|
|
869
|
+
seek: list[str] = []
|
|
870
|
+
if f0 > 0 and abs(base_fps - fps) <= 0.01:
|
|
871
|
+
seek = ["-ss", f"{max(0.0, (f0 - 0.25) / fps):.6f}"]
|
|
872
|
+
elif f0 > 0:
|
|
873
|
+
filters.append(f"select='gte(n\\,{f0})'")
|
|
874
|
+
filters.append(f"format={pix_fmt}")
|
|
875
|
+
return _ff() + seek + ["-i", str(base), "-map", "0:v:0", "-vf", ",".join(filters), "-fps_mode", "passthrough",
|
|
876
|
+
"-frames:v", str(n), "-an", "-f", "rawvideo", "-pix_fmt", pix_fmt, "-"]
|
|
877
|
+
|
|
878
|
+
|
|
879
|
+
def matte_cmd(path: Path, size: tuple[int, int] | None = None, full_size: tuple[int, int] | None = None) -> list[str]:
|
|
880
|
+
vf = ["format=gray"]
|
|
881
|
+
if size is not None and full_size is not None and tuple(size) != tuple(full_size):
|
|
882
|
+
vf.insert(0, f"scale={size[0]}:{size[1]}:flags=area")
|
|
883
|
+
return _ff() + ["-i", str(path), "-map", "0:v:0", "-vf", ",".join(vf), "-fps_mode", "passthrough",
|
|
884
|
+
"-an", "-f", "rawvideo", "-pix_fmt", "gray", "-"]
|
|
885
|
+
|
|
886
|
+
|
|
887
|
+
def _encoder(tmp: Path, size: tuple[int, int], fps: int) -> subprocess.Popen:
|
|
888
|
+
W, H = size
|
|
889
|
+
return subprocess.Popen(_ff() + ["-y", "-f", "rawvideo", "-pix_fmt", "gray", "-s", f"{W}x{H}", "-framerate", str(fps),
|
|
890
|
+
"-i", "-", "-c:v", "ffv1", "-level", "3", "-slices", "4", "-g", "1",
|
|
891
|
+
"-pix_fmt", "gray", "-f", "matroska", str(tmp)],
|
|
892
|
+
stdin=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
893
|
+
|
|
894
|
+
|
|
895
|
+
def count_frames(path: Path) -> int:
|
|
896
|
+
proc = subprocess.run(["ffprobe", "-v", "error", "-select_streams", "v:0", "-count_packets", "-show_entries",
|
|
897
|
+
"stream=nb_read_packets", "-of", "csv=p=0", str(path)], capture_output=True, text=True)
|
|
898
|
+
try:
|
|
899
|
+
return int(proc.stdout.strip().split(",")[0])
|
|
900
|
+
except ValueError:
|
|
901
|
+
return -1
|
|
902
|
+
|
|
903
|
+
|
|
904
|
+
def _small(frame: np.ndarray) -> np.ndarray:
|
|
905
|
+
cv2 = _cv2()
|
|
906
|
+
h, w = frame.shape[:2]
|
|
907
|
+
k = SYNC_SIZE / min(h, w)
|
|
908
|
+
size = (max(8, int(round(w * k))), max(8, int(round(h * k))))
|
|
909
|
+
g = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY) if frame.ndim == 3 else frame
|
|
910
|
+
return cv2.resize(g, size, interpolation=cv2.INTER_AREA).astype(np.float32)
|
|
911
|
+
|
|
912
|
+
|
|
913
|
+
def _read_all(cmd: list[str], size: tuple[int, int], what: str) -> list[np.ndarray]:
|
|
914
|
+
W, H = size
|
|
915
|
+
proc = subprocess.run(cmd, capture_output=True)
|
|
916
|
+
n = W * H * 3
|
|
917
|
+
data = proc.stdout
|
|
918
|
+
if proc.returncode != 0 and not data:
|
|
919
|
+
raise MatteError(f"o ffmpeg falhou ao {what}: {proc.stderr.decode('utf-8', 'replace').strip()[-300:]}")
|
|
920
|
+
return [np.frombuffer(data[i:i + n], np.uint8).reshape(H, W, 3) for i in range(0, len(data) - n + 1, n)]
|
|
921
|
+
|
|
922
|
+
|
|
923
|
+
def sync_check(job: RangeJob, got_head: list[np.ndarray], got_tail: list[np.ndarray]) -> tuple[bool, dict]:
|
|
924
|
+
"""Do the first and last frames read from the source match cut.mp4 at the same positions?"""
|
|
925
|
+
out = {}
|
|
926
|
+
ok = True
|
|
927
|
+
for name, frames, start in (("inicio", got_head, job.f0), ("fim", got_tail, job.f1 - len(got_tail))):
|
|
928
|
+
if not frames:
|
|
929
|
+
continue
|
|
930
|
+
lo = max(0, start - 2)
|
|
931
|
+
n = len(frames) + 4 + (start - 2 - lo)
|
|
932
|
+
ref = _read_all(base_cmd(job.base, job.base_pre, job.base_fps, job.fps, lo, n), job.size,
|
|
933
|
+
"ler o corte para conferir o recorte")
|
|
934
|
+
ref_s = [_small(f) for f in ref]
|
|
935
|
+
got_s = [_small(f) for f in frames]
|
|
936
|
+
mads = {}
|
|
937
|
+
for off in (-2, -1, 0, 1, 2):
|
|
938
|
+
diffs = []
|
|
939
|
+
for i, g in enumerate(got_s):
|
|
940
|
+
j = start + off + i - lo
|
|
941
|
+
if 0 <= j < len(ref_s) and ref_s[j].shape == g.shape:
|
|
942
|
+
diffs.append(float(np.abs(ref_s[j] - g).mean()))
|
|
943
|
+
if diffs:
|
|
944
|
+
mads[off] = sum(diffs) / len(diffs)
|
|
945
|
+
if 0 not in mads:
|
|
946
|
+
ok = False
|
|
947
|
+
out[name] = {"mad": None}
|
|
948
|
+
continue
|
|
949
|
+
best = min(mads.values())
|
|
950
|
+
good = mads[0] <= SYNC_MAX_MAD and mads[0] <= best * 1.3 + 0.1
|
|
951
|
+
out[name] = {"mad": round(mads[0], 3), "best_offset": min(mads, key=mads.get), "ok": good}
|
|
952
|
+
ok = ok and good
|
|
953
|
+
return ok, out
|
|
954
|
+
|
|
955
|
+
|
|
956
|
+
# ================================================================== computing one range
|
|
957
|
+
|
|
958
|
+
def _progress_line(msg: str, progress: Callable[[str], None] | None) -> None:
|
|
959
|
+
if progress:
|
|
960
|
+
progress(msg)
|
|
961
|
+
|
|
962
|
+
|
|
963
|
+
def compute_job(job: RangeJob, engine, progress: Callable[[str], None] | None = None,
|
|
964
|
+
on_frames: Callable[[int], None] | None = None) -> dict:
|
|
965
|
+
"""Run the engine over one range and write its matte file. Returns the meta."""
|
|
966
|
+
W, H = job.size
|
|
967
|
+
job.path.parent.mkdir(parents=True, exist_ok=True)
|
|
968
|
+
tmp = job.path.with_name(".tmp-" + job.path.name)
|
|
969
|
+
started = time.time()
|
|
970
|
+
method = job.method
|
|
971
|
+
if method == "fonte":
|
|
972
|
+
cmd = source_cmd(job, job.source_hw)
|
|
973
|
+
else:
|
|
974
|
+
cmd = base_cmd(job.base, job.base_pre, job.base_fps, job.fps, job.f0, job.frames)
|
|
975
|
+
reader = FrameReader(cmd, (W, H), 3, "ler o vídeo")
|
|
976
|
+
enc = _encoder(tmp, (W, H), job.fps)
|
|
977
|
+
rows: list = []
|
|
978
|
+
head: list[np.ndarray] = []
|
|
979
|
+
tail: deque = deque(maxlen=SYNC_FRAMES)
|
|
980
|
+
pending: deque = deque()
|
|
981
|
+
written = 0
|
|
982
|
+
engine.reset()
|
|
983
|
+
pool = ThreadPoolExecutor(max_workers=2)
|
|
984
|
+
err = ""
|
|
985
|
+
|
|
986
|
+
def write_ready(block: bool) -> None:
|
|
987
|
+
nonlocal written
|
|
988
|
+
while pending and (block or pending[0].done()):
|
|
989
|
+
alpha = pending.popleft().result()
|
|
990
|
+
enc.stdin.write((np.clip(alpha, 0.0, 1.0) * 255.0 + 0.5).astype(np.uint8).tobytes())
|
|
991
|
+
written += 1
|
|
992
|
+
if on_frames:
|
|
993
|
+
on_frames(1)
|
|
994
|
+
|
|
995
|
+
st = SETTINGS["start"]
|
|
996
|
+
settle = getattr(engine, "settle_start", None)
|
|
997
|
+
n_start = min(st["frames"], job.frames // 3) if settle is not None and job.frames >= st["min_range"] else 0
|
|
998
|
+
first: list[tuple] = [] # (frame, temporal matte, probability, gray) of the first n_start frames
|
|
999
|
+
|
|
1000
|
+
def emit(frame: np.ndarray, coarse: np.ndarray) -> None:
|
|
1001
|
+
rows.append(frame_stats(coarse))
|
|
1002
|
+
pending.append(pool.submit(engine.refine, frame, coarse))
|
|
1003
|
+
write_ready(len(pending) >= 3)
|
|
1004
|
+
|
|
1005
|
+
try:
|
|
1006
|
+
for i in range(job.frames):
|
|
1007
|
+
frame = reader.read()
|
|
1008
|
+
if frame is None:
|
|
1009
|
+
break
|
|
1010
|
+
if method == "fonte":
|
|
1011
|
+
if len(head) < SYNC_FRAMES:
|
|
1012
|
+
head.append(frame)
|
|
1013
|
+
tail.append(frame)
|
|
1014
|
+
if not n_start:
|
|
1015
|
+
emit(frame, engine.coarse(frame))
|
|
1016
|
+
continue
|
|
1017
|
+
a, p, gray = engine.coarse_parts(frame)
|
|
1018
|
+
if i < n_start:
|
|
1019
|
+
first.append((frame, a, p, gray))
|
|
1020
|
+
continue
|
|
1021
|
+
if i == n_start:
|
|
1022
|
+
fixed = settle([(fa, fp, fg) for _, fa, fp, fg in first], (a, gray))
|
|
1023
|
+
for (f_rgb, _, fp, _), fa in zip(first, fixed):
|
|
1024
|
+
emit(f_rgb, engine.finish(fa, fp))
|
|
1025
|
+
first = []
|
|
1026
|
+
emit(frame, engine.finish(a, p))
|
|
1027
|
+
for f_rgb, fa, fp, _ in first: # a range that ended before the anchor frame
|
|
1028
|
+
emit(f_rgb, engine.finish(fa, fp))
|
|
1029
|
+
write_ready(True)
|
|
1030
|
+
except BrokenPipeError:
|
|
1031
|
+
pass
|
|
1032
|
+
finally:
|
|
1033
|
+
pool.shutdown(wait=True)
|
|
1034
|
+
err = reader.close()
|
|
1035
|
+
try:
|
|
1036
|
+
enc.stdin.close()
|
|
1037
|
+
except OSError:
|
|
1038
|
+
pass
|
|
1039
|
+
enc_err = enc.stderr.read().decode("utf-8", "replace") if enc.stderr else ""
|
|
1040
|
+
code = enc.wait()
|
|
1041
|
+
if code != 0:
|
|
1042
|
+
tmp.unlink(missing_ok=True)
|
|
1043
|
+
raise MatteError(f"o ffmpeg falhou ao gravar o recorte do trecho {job.rid}: {enc_err.strip()[-300:]}")
|
|
1044
|
+
if written == 0:
|
|
1045
|
+
tmp.unlink(missing_ok=True)
|
|
1046
|
+
if method == "fonte":
|
|
1047
|
+
return {"retry": "nenhum quadro lido do vídeo bruto", "written": 0}
|
|
1048
|
+
raise MatteError(f"não consegui ler os quadros do trecho {job.rid} ({err[-300:] or 'vídeo vazio'})")
|
|
1049
|
+
sync = None
|
|
1050
|
+
if method == "fonte":
|
|
1051
|
+
ok, sync = sync_check(job, head, list(tail))
|
|
1052
|
+
if not ok or written != job.frames:
|
|
1053
|
+
tmp.unlink(missing_ok=True)
|
|
1054
|
+
why = "os quadros do vídeo bruto não batem com o corte" if not ok else "o vídeo bruto terminou antes"
|
|
1055
|
+
return {"retry": why, "sync": sync, "written": written}
|
|
1056
|
+
if written < job.frames:
|
|
1057
|
+
# the base ended early: hold the last matte so the frame count stays exact
|
|
1058
|
+
tmp.unlink(missing_ok=True)
|
|
1059
|
+
raise MatteError(f"o trecho {job.rid} tem {written} quadros e o corte espera {job.frames}. Renderize o corte de novo")
|
|
1060
|
+
got = count_frames(tmp)
|
|
1061
|
+
if got != job.frames:
|
|
1062
|
+
tmp.unlink(missing_ok=True)
|
|
1063
|
+
raise MatteError(f"o recorte do trecho {job.rid} saiu com {got} quadros, esperava {job.frames}")
|
|
1064
|
+
os.replace(tmp, job.path)
|
|
1065
|
+
seconds = time.time() - started
|
|
1066
|
+
meta = {"key": job.key, "engine": engine_version(), "provider": getattr(engine, "provider", "?"),
|
|
1067
|
+
"id": job.rid, "method": method, "source": str(job.source) if job.source else None,
|
|
1068
|
+
"base": job.base.name if job.base else None, "note": job.note, "frames": job.frames,
|
|
1069
|
+
"size": list(job.size), "fps": job.fps, "stats": summarize_stats(rows), "sync": sync,
|
|
1070
|
+
"seconds": round(seconds, 2), "created_at": now_iso()}
|
|
1071
|
+
save_json(job.meta_path, meta)
|
|
1072
|
+
job.seconds = seconds
|
|
1073
|
+
return meta
|
|
1074
|
+
|
|
1075
|
+
|
|
1076
|
+
def _load_cached(job: RangeJob) -> dict | None:
|
|
1077
|
+
meta = load_json(job.meta_path) if job.meta_path.exists() else None
|
|
1078
|
+
if (isinstance(meta, dict) and meta.get("key") == job.key and job.path.exists()
|
|
1079
|
+
and job.path.stat().st_size > 0 and int(meta.get("frames") or -1) == job.frames):
|
|
1080
|
+
return meta
|
|
1081
|
+
return None
|
|
1082
|
+
|
|
1083
|
+
|
|
1084
|
+
@dataclass
|
|
1085
|
+
class MatteSet:
|
|
1086
|
+
jobs: list[RangeJob]
|
|
1087
|
+
size: tuple[int, int]
|
|
1088
|
+
fps: int
|
|
1089
|
+
seconds: float = 0.0
|
|
1090
|
+
computed: int = 0
|
|
1091
|
+
warnings: list[str] = field(default_factory=list)
|
|
1092
|
+
provider: str = ""
|
|
1093
|
+
|
|
1094
|
+
@property
|
|
1095
|
+
def cached(self) -> int:
|
|
1096
|
+
return sum(1 for j in self.jobs if j.cached)
|
|
1097
|
+
|
|
1098
|
+
def report(self, rel: Callable[[Path], str] = str) -> dict:
|
|
1099
|
+
return {"engine": engine_version(), "model": MODEL_NAME, "provider": self.provider or None,
|
|
1100
|
+
"seconds": round(self.seconds, 2), "cached": self.cached, "computed": self.computed,
|
|
1101
|
+
"ranges": [{"id": j.rid, "start_frame": j.f0, "end_frame": j.f1, "method": j.method,
|
|
1102
|
+
"note": j.note or None, "file": rel(j.path), "cached": j.cached,
|
|
1103
|
+
"seconds": j.meta.get("seconds"), "stats": j.meta.get("stats")} for j in self.jobs]}
|
|
1104
|
+
|
|
1105
|
+
def frames(self, size: tuple[int, int] | None = None) -> Iterator[np.ndarray]:
|
|
1106
|
+
"""Every matte frame in order (uint8 HxW), optionally scaled to `size`."""
|
|
1107
|
+
W, H = size or self.size
|
|
1108
|
+
for job in self.jobs:
|
|
1109
|
+
reader = FrameReader(matte_cmd(job.path, (W, H), self.size), (W, H), 1, "ler o recorte")
|
|
1110
|
+
got = 0
|
|
1111
|
+
last = None
|
|
1112
|
+
try:
|
|
1113
|
+
for _ in range(job.frames):
|
|
1114
|
+
fr = reader.read()
|
|
1115
|
+
if fr is None:
|
|
1116
|
+
break
|
|
1117
|
+
last = fr
|
|
1118
|
+
got += 1
|
|
1119
|
+
yield fr
|
|
1120
|
+
finally:
|
|
1121
|
+
err = reader.close()
|
|
1122
|
+
if got < job.frames:
|
|
1123
|
+
if last is None:
|
|
1124
|
+
raise MatteError(f"não consegui ler o recorte do trecho {job.rid} ({err[-200:]})")
|
|
1125
|
+
for _ in range(job.frames - got):
|
|
1126
|
+
yield last
|
|
1127
|
+
|
|
1128
|
+
|
|
1129
|
+
def ensure(edit_dir: Path, base: Path, bounds: list[tuple[int, int]], fps: int, size: tuple[int, int],
|
|
1130
|
+
timeline: dict | None, progress: Callable[[str], None] | None = print, force: bool = False,
|
|
1131
|
+
base_fps: float | None = None, base_size: tuple[int, int] | None = None) -> MatteSet:
|
|
1132
|
+
"""The matte of every range, computed now when it is not in the cache."""
|
|
1133
|
+
require_binary("ffmpeg")
|
|
1134
|
+
jobs = plan_jobs(edit_dir, base, bounds, fps, size, timeline, base_fps, base_size)
|
|
1135
|
+
mset = MatteSet(jobs=jobs, size=tuple(size), fps=fps)
|
|
1136
|
+
todo = []
|
|
1137
|
+
for job in jobs:
|
|
1138
|
+
meta = None if force else _load_cached(job)
|
|
1139
|
+
if meta is not None:
|
|
1140
|
+
job.cached, job.meta = True, meta
|
|
1141
|
+
mset.provider = mset.provider or str(meta.get("provider") or "")
|
|
1142
|
+
else:
|
|
1143
|
+
todo.append(job)
|
|
1144
|
+
if not todo:
|
|
1145
|
+
return mset
|
|
1146
|
+
default_engine_ready()
|
|
1147
|
+
total_frames = sum(j.frames for j in todo)
|
|
1148
|
+
video_s = total_frames / fps
|
|
1149
|
+
est = video_s * SECONDS_PER_SECOND
|
|
1150
|
+
cached_note = f", {len(jobs) - len(todo)} já prontos" if len(todo) < len(jobs) else ""
|
|
1151
|
+
_progress_line(f"recorte da pessoa: {len(todo)} trecho(s), {video_s:.1f}s de vídeo (leva uns "
|
|
1152
|
+
f"{_duration_words(est)}{cached_note})", progress)
|
|
1153
|
+
started = time.time()
|
|
1154
|
+
done_frames = [0]
|
|
1155
|
+
last_print = [started]
|
|
1156
|
+
|
|
1157
|
+
def tick(n: int) -> None:
|
|
1158
|
+
done_frames[0] += n
|
|
1159
|
+
now = time.time()
|
|
1160
|
+
if now - last_print[0] >= PROGRESS_EVERY and done_frames[0] < total_frames:
|
|
1161
|
+
last_print[0] = now
|
|
1162
|
+
pct = 100.0 * done_frames[0] / max(1, total_frames)
|
|
1163
|
+
rate = done_frames[0] / max(1e-6, now - started)
|
|
1164
|
+
left = (total_frames - done_frames[0]) / max(1e-6, rate)
|
|
1165
|
+
_progress_line(f"recorte da pessoa: {pct:.0f}% (falta uns {_duration_words(left)})", progress)
|
|
1166
|
+
|
|
1167
|
+
engine = make_engine(tuple(size))
|
|
1168
|
+
mset.provider = str(getattr(engine, "provider", "") or "")
|
|
1169
|
+
try:
|
|
1170
|
+
for job in todo:
|
|
1171
|
+
meta = compute_job(job, engine, progress, tick)
|
|
1172
|
+
if "retry" in meta:
|
|
1173
|
+
why = meta["retry"]
|
|
1174
|
+
# back to cut.mp4 for this range: frames already counted are counted again
|
|
1175
|
+
done_frames[0] -= min(done_frames[0], int(meta.get("written") or 0))
|
|
1176
|
+
job.method, job.note, job.source = "corte", why, None
|
|
1177
|
+
job.key = _short({"v": engine_version(), "s": SETTINGS, "how": "corte", "base": job.base.name,
|
|
1178
|
+
"sig": file_signature(job.base), "f0": job.f0, "f1": job.f1, "size": list(size),
|
|
1179
|
+
"fps": fps, "pre": job.base_pre})
|
|
1180
|
+
job.path = job.path.with_name(f"m_{job.key}.mkv")
|
|
1181
|
+
job.meta_path = job.path.with_suffix(".json")
|
|
1182
|
+
cached = None if force else _load_cached(job)
|
|
1183
|
+
if cached is not None:
|
|
1184
|
+
job.cached, job.meta = True, cached
|
|
1185
|
+
continue
|
|
1186
|
+
meta = compute_job(job, engine, progress, tick)
|
|
1187
|
+
job.meta = meta
|
|
1188
|
+
mset.computed += 1
|
|
1189
|
+
finally:
|
|
1190
|
+
try:
|
|
1191
|
+
engine.close()
|
|
1192
|
+
except Exception: # noqa: BLE001
|
|
1193
|
+
pass
|
|
1194
|
+
mset.seconds = time.time() - started
|
|
1195
|
+
notes = sorted({j.note for j in jobs if j.method == "corte" and j.note and not j.cached})
|
|
1196
|
+
if notes and timeline:
|
|
1197
|
+
mset.warnings.append("recorte feito a partir do cut.mp4 em vez do vídeo bruto (" + "; ".join(notes) + ")")
|
|
1198
|
+
_progress_line(f"recorte da pessoa pronto: {len(jobs)} trecho(s) ({mset.cached} do cache) em {mset.seconds:.1f}s "
|
|
1199
|
+
f"({mset.seconds / max(video_s, 1e-6):.1f}x a duração)", progress)
|
|
1200
|
+
return mset
|
|
1201
|
+
|
|
1202
|
+
|
|
1203
|
+
def _duration_words(seconds: float) -> str:
|
|
1204
|
+
if seconds < 90:
|
|
1205
|
+
return f"{max(1, int(round(seconds)))} s"
|
|
1206
|
+
return f"{int(round(seconds / 60.0))} min"
|
|
1207
|
+
|
|
1208
|
+
|
|
1209
|
+
def estimate_seconds(duration: float) -> float:
|
|
1210
|
+
return duration * SECONDS_PER_SECOND
|
|
1211
|
+
|
|
1212
|
+
|
|
1213
|
+
# ================================================================== CLI
|
|
1214
|
+
|
|
1215
|
+
def _segments_bounds(timeline: dict | None, nf: int, fps: int) -> list[tuple[int, int]]:
|
|
1216
|
+
starts = []
|
|
1217
|
+
for s in (timeline or {}).get("segments") or []:
|
|
1218
|
+
try:
|
|
1219
|
+
starts.append(int(round(float(s["out_start"]) * fps + 1e-6)))
|
|
1220
|
+
except (KeyError, TypeError, ValueError):
|
|
1221
|
+
continue
|
|
1222
|
+
cuts = sorted({f for f in starts if 0 < f < nf})
|
|
1223
|
+
b = [0] + cuts + [nf]
|
|
1224
|
+
return [(x, y) for x, y in zip(b, b[1:]) if y > x]
|
|
1225
|
+
|
|
1226
|
+
|
|
1227
|
+
def _base_facts(base: Path) -> tuple[float, tuple[int, int], int]:
|
|
1228
|
+
info = probe_video(base)
|
|
1229
|
+
fps = float(info.get("fps") or 30.0)
|
|
1230
|
+
size = (int(info.get("display_width") or info.get("width") or 0), int(info.get("display_height") or info.get("height") or 0))
|
|
1231
|
+
frames = count_frames(base)
|
|
1232
|
+
if frames <= 0:
|
|
1233
|
+
frames = int(round(float(info.get("duration") or 0) * fps))
|
|
1234
|
+
return fps, size, frames
|
|
1235
|
+
|
|
1236
|
+
|
|
1237
|
+
def run_edit(edit_dir: Path, video: Path | None, force: bool) -> tuple[MatteSet, dict]:
|
|
1238
|
+
edit_dir = edit_dir.resolve()
|
|
1239
|
+
base = (video or edit_dir / "cut.mp4").resolve()
|
|
1240
|
+
if not base.exists():
|
|
1241
|
+
raise MatteError(f"vídeo não encontrado: {base}", 2)
|
|
1242
|
+
timeline = load_json(edit_dir / "cut_timeline.json") if base == edit_dir / "cut.mp4" else None
|
|
1243
|
+
base_fps, base_size, nf = _base_facts(base)
|
|
1244
|
+
from common import output_fps_for # noqa: PLC0415
|
|
1245
|
+
fps = int(round(float((timeline or {}).get("fps") or output_fps_for(base_fps))))
|
|
1246
|
+
size = (int((timeline or {}).get("width") or base_size[0]), int((timeline or {}).get("height") or base_size[1]))
|
|
1247
|
+
size = (size[0] - size[0] % 2, size[1] - size[1] % 2)
|
|
1248
|
+
if abs(base_fps - fps) > 0.01:
|
|
1249
|
+
nf = int(round(nf / base_fps * fps))
|
|
1250
|
+
bounds = _segments_bounds(timeline, nf, fps)
|
|
1251
|
+
mset = ensure(edit_dir, base, bounds, fps, size, timeline, progress=lambda m: print(m, flush=True), force=force,
|
|
1252
|
+
base_fps=base_fps, base_size=base_size)
|
|
1253
|
+
return mset, {"base": base, "frames": nf, "fps": fps}
|
|
1254
|
+
|
|
1255
|
+
|
|
1256
|
+
def image_matte(src: Path, out: Path) -> dict:
|
|
1257
|
+
"""One photo: the cut-out as a PNG with transparency (no temporal filter)."""
|
|
1258
|
+
cv2 = _cv2()
|
|
1259
|
+
bgr = cv2.imread(str(src), cv2.IMREAD_COLOR)
|
|
1260
|
+
if bgr is None:
|
|
1261
|
+
raise MatteError(f"não consegui abrir a imagem {src}", 2)
|
|
1262
|
+
h, w = bgr.shape[:2]
|
|
1263
|
+
rgb = np.ascontiguousarray(bgr[..., ::-1])
|
|
1264
|
+
rgb = np.ascontiguousarray(rgb[:h - h % 2, :w - w % 2])
|
|
1265
|
+
engine = make_engine((rgb.shape[1], rgb.shape[0]))
|
|
1266
|
+
try:
|
|
1267
|
+
alpha = engine.refine(rgb, engine.coarse(rgb))
|
|
1268
|
+
finally:
|
|
1269
|
+
engine.close()
|
|
1270
|
+
a8 = (np.clip(alpha, 0, 1) * 255 + 0.5).astype(np.uint8)
|
|
1271
|
+
rgba = np.dstack([rgb[..., ::-1], a8])
|
|
1272
|
+
out.parent.mkdir(parents=True, exist_ok=True)
|
|
1273
|
+
cv2.imwrite(str(out), rgba)
|
|
1274
|
+
return {"file": str(out), "width": rgb.shape[1], "height": rgb.shape[0],
|
|
1275
|
+
"person_share": round(float((alpha > 0.5).mean()), 4)}
|
|
1276
|
+
|
|
1277
|
+
|
|
1278
|
+
def main(argv: list[str] | None = None) -> int:
|
|
1279
|
+
for stream in (sys.stdout, sys.stderr):
|
|
1280
|
+
try:
|
|
1281
|
+
stream.reconfigure(encoding="utf-8", errors="replace") # type: ignore[attr-defined]
|
|
1282
|
+
except (AttributeError, ValueError, OSError):
|
|
1283
|
+
pass
|
|
1284
|
+
ap = argparse.ArgumentParser(description="Recorta a pessoa de cada trecho do corte (tela verde) e guarda o "
|
|
1285
|
+
"recorte para o compose_final.py. Leva uns 3 minutos por minuto de vídeo.")
|
|
1286
|
+
ap.add_argument("--edit-dir", help="pasta edit da sessão")
|
|
1287
|
+
ap.add_argument("--video", help="vídeo a recortar (padrão: edit/cut.mp4)")
|
|
1288
|
+
ap.add_argument("--force", action="store_true", help="recorta de novo mesmo com o recorte guardado")
|
|
1289
|
+
ap.add_argument("--json", action="store_true", help="imprime o resumo em JSON")
|
|
1290
|
+
ap.add_argument("--check", action="store_true", help="só confere se o recorte funciona neste computador")
|
|
1291
|
+
ap.add_argument("--image", help="recorta uma foto (teste rápido)")
|
|
1292
|
+
ap.add_argument("-o", "--output", help="com --image: PNG de saída")
|
|
1293
|
+
args = ap.parse_args(argv)
|
|
1294
|
+
try:
|
|
1295
|
+
if args.check:
|
|
1296
|
+
info = check_engine()
|
|
1297
|
+
seg = OnnxSegmenter(sessions=1)
|
|
1298
|
+
t0 = time.perf_counter()
|
|
1299
|
+
y = seg.raw(np.zeros((64, 64, 3), np.uint8))
|
|
1300
|
+
ms = (time.perf_counter() - t0) * 1000
|
|
1301
|
+
seg.close()
|
|
1302
|
+
print(f"ok: recorte da pessoa pronto (onnxruntime {info['onnxruntime']}, {seg.provider}, saída "
|
|
1303
|
+
f"{tuple(y.shape)}, {ms:.0f} ms no primeiro quadro)")
|
|
1304
|
+
return 0
|
|
1305
|
+
if args.image:
|
|
1306
|
+
if not args.output:
|
|
1307
|
+
print("erro: use -o arquivo.png junto com --image", file=sys.stderr)
|
|
1308
|
+
return 2
|
|
1309
|
+
res = image_matte(Path(args.image).expanduser(), Path(args.output).expanduser())
|
|
1310
|
+
print(f"ok: {res['file']} ({res['width']}x{res['height']}, pessoa em {res['person_share'] * 100:.0f}% da imagem)")
|
|
1311
|
+
return 0
|
|
1312
|
+
if not args.edit_dir:
|
|
1313
|
+
print("erro: diga a pasta com --edit-dir (ou use --check)", file=sys.stderr)
|
|
1314
|
+
return 2
|
|
1315
|
+
edit = Path(args.edit_dir).expanduser()
|
|
1316
|
+
if not edit.is_dir():
|
|
1317
|
+
print(f"erro: pasta de edição não encontrada: {edit}", file=sys.stderr)
|
|
1318
|
+
return 2
|
|
1319
|
+
mset, facts = run_edit(edit, Path(args.video).expanduser() if args.video else None, args.force)
|
|
1320
|
+
rep = mset.report(lambda p: str(p.relative_to(edit.resolve())) if p.is_relative_to(edit.resolve()) else str(p))
|
|
1321
|
+
for w in mset.warnings:
|
|
1322
|
+
print(f"aviso: {w}")
|
|
1323
|
+
if args.json:
|
|
1324
|
+
print(json.dumps(rep, ensure_ascii=False, indent=2))
|
|
1325
|
+
return 0
|
|
1326
|
+
for j in mset.jobs:
|
|
1327
|
+
st = j.meta.get("stats") or {}
|
|
1328
|
+
who = "sem pessoa" if not st.get("person_frames") else f"cabeça em {st.get('top', 0):.2f} da altura"
|
|
1329
|
+
print(f" {j.rid:>5} {fmt_time(j.f0 / facts['fps'])}-{fmt_time(j.f1 / facts['fps'])} "
|
|
1330
|
+
f"{'cache' if j.cached else j.method:>6} {who}")
|
|
1331
|
+
print(f"ok: {len(mset.jobs)} trecho(s) recortados em {edit.resolve() / CACHE_DIR}")
|
|
1332
|
+
return 0
|
|
1333
|
+
except MatteError as e:
|
|
1334
|
+
print(f"erro: {e}", file=sys.stderr)
|
|
1335
|
+
return e.code
|
|
1336
|
+
except KeyboardInterrupt:
|
|
1337
|
+
print("interrompido", file=sys.stderr)
|
|
1338
|
+
return 130
|
|
1339
|
+
|
|
1340
|
+
|
|
1341
|
+
if __name__ == "__main__":
|
|
1342
|
+
sys.exit(main())
|