openTARS / app /tars_modelinhos.py

Mirror of openTARS 5.2 Ranger (revision 2026110101)

2695555 verified 6 days ago

9.14 kB

"""5.2 (decisão do Enzo, 25/10): os modelinhos de foto, todos baixados obrigatoriamente e rodando na CPU (ONNX Runtime):
tirar o fundo (BiRefNet lite), aumentar a resolução (Real-ESRGAN general x4v3), selecionar por clique (MobileSAM) e
remover um objeto (MI-GAN, com a máscara do selecionar). Cada arquivo tem endereço com o commit fixo e o SHA-256
conferido; o que não confere é apagado. Sem Qt.
Medido em 25/10 no PC do Enzo (16 threads, só CPU): tirar o fundo ~9 s por foto (3 s de carga na primeira), selecionar
~0,5 s, remover ~1 s, aumentar ~0,5 s a cada pedaço de 256 px."""
import hashlib
import os
import threading
import urllib.request
from pathlib import Path

import tars_plataforma as plataforma

_HF = "https://huggingface.co/"
# nome: (arquivo, endereço com o commit fixo, SHA-256, bytes, licença)
CATALOGO = {
    "fundo": ("birefnet_lite.onnx",
              _HF + "onnx-community/BiRefNet_lite-ONNX/resolve/de15b22ba131738a16dff04aab8bdf8dc32e3ac1/onnx/model.onnx",
              "5600024376f572a557870a5eb0afb1e5961636bef4e1e22132025467d0f03333", 224005088, "MIT"),
    "aumentar": ("realesr_general_x4v3.onnx",
                 _HF + "vlad-grig/realesr-general-x4v3-onnx/resolve/0e95d93e454cf6b5de9bdf5a1c9e7ec340d38231/onnx/model.onnx",
                 "7b10e80021ec3e0ff35017ae5022791879d93a799b2e201e1010a37510202266", 4866420, "BSD-3-Clause"),
    "remover": ("migan_pipeline_v2.onnx",
                _HF + "andraniksargsyan/migan/resolve/406830d0fa60666da0071c342ad2fbc8f30c5c64/migan_pipeline_v2.onnx",
                "6f1f3530a1a2324b19752018ce756088b07973cda8d7d890034ace5c8a48c40b", 28079181, "MIT"),
    "selecionar_codificador": ("mobilesam_encoder.onnx",
                               _HF + "Acly/MobileSAM/resolve/0d3b403339b4674a82493d5e97964dd78089ddc8/mobile_sam_image_encoder.onnx",
                               "580f5fb648ea1062c0aabc26217aed56921985f03f0cbbd852bba81d760cc749", 28157093, "MIT"),
    "selecionar_decodificador": ("mobilesam_decoder.onnx",
                                 _HF + "Acly/MobileSAM/resolve/0d3b403339b4674a82493d5e97964dd78089ddc8/sam_mask_decoder_single.onnx",
                                 "93915fc7c993ab9d59ab8c9ccd3bce37f7509c81ab4150a74abd4d2abbd8570d", 16501323, "MIT"),
}
LADO_MAXIMO_AUMENTADO = 4096        # aumentar nunca passa disso no lado maior (memória e tempo na CPU)


def pasta():
    return Path(os.environ.get("TARS_PASTA_MODELINHOS") or plataforma.pastas()["dados"] / "modelinhos")


def caminho(nome):
    return pasta() / CATALOGO[nome][0]


def faltando():
    """Os modelinhos que ainda não estão no disco (pelo tamanho; o hash foi conferido ao baixar). Os testes do app
    (TARS_SEM_MODELINHOS=1) não baixam ~300 MB da internet."""
    if os.environ.get("TARS_SEM_MODELINHOS") == "1":
        return []
    return [n for n, (_, _, _, tam, _) in CATALOGO.items() if not (caminho(n).is_file() and caminho(n).stat().st_size == tam)]


def bytes_faltando():
    return sum(CATALOGO[n][3] for n in faltando())


def baixar(nome, ao_progredir=None, cancelado=None, abrir=None):
    """Baixa um modelinho para um .part, confere o SHA-256 e só então põe no lugar. (ok, erro)."""
    arquivo, url, sha, tamanho, _ = CATALOGO[nome]
    destino = pasta() / arquivo
    destino.parent.mkdir(parents=True, exist_ok=True)
    parcial = destino.with_name(arquivo + ".part")
    soma, feito = hashlib.sha256(), 0
    try:
        with (abrir or urllib.request.urlopen)(urllib.request.Request(url, headers={"User-Agent": "openTARS"}), timeout=60) as r, \
                open(parcial, "wb") as f:
            while True:
                if cancelado and cancelado():
                    raise InterruptedError("cancelado")
                bloco = r.read(1 << 20)
                if not bloco:
                    break
                f.write(bloco)
                soma.update(bloco)
                feito += len(bloco)
                if ao_progredir:
                    ao_progredir(feito, tamanho)
    except InterruptedError:
        parcial.unlink(missing_ok=True)
        return False, "cancelado"
    except OSError as e:
        parcial.unlink(missing_ok=True)
        return False, str(e)
    if soma.hexdigest() != sha:
        parcial.unlink(missing_ok=True)
        return False, f"SHA-256 diferente em {arquivo}"
    os.replace(parcial, destino)
    return True, ""


def baixar_todos(ao_progredir=None, cancelado=None, abrir=None):
    """Baixa o que falta, um por um. ao_progredir(feito, total) é do conjunto. (ok, erro)."""
    nomes = faltando()
    total, antes = sum(CATALOGO[n][3] for n in nomes), 0
    for n in nomes:
        ok, erro = baixar(n, (lambda f, _t, a=antes: ao_progredir(a + f, total)) if ao_progredir else None, cancelado, abrir)
        if not ok:
            return False, erro
        antes += CATALOGO[n][3]
    return True, ""


# ------------------------------------------------------------
# Os modelos (sessões abertas uma vez e guardadas)
# ------------------------------------------------------------
_SESSOES, _TRAVA = {}, threading.Lock()


def _sessao(nome):
    with _TRAVA:
        if nome not in _SESSOES:
            import onnxruntime
            if not caminho(nome).is_file():
                raise FileNotFoundError(f"modelinho {nome} não baixado")
            _SESSOES[nome] = onnxruntime.InferenceSession(str(caminho(nome)), providers=["CPUExecutionProvider"])
        return _SESSOES[nome]


def tirar_fundo(imagem):
    """A foto (PIL) com o fundo transparente (RGBA)."""
    import numpy as np
    from PIL import Image
    rgb = imagem.convert("RGB")
    x = np.asarray(rgb.resize((1024, 1024), Image.BILINEAR), dtype=np.float32) / 255.0
    x = (x - (0.485, 0.456, 0.406)) / (0.229, 0.224, 0.225)
    saida = _sessao("fundo").run(None, {"input_image": x.transpose(2, 0, 1)[None].astype(np.float32)})[0][0, 0]
    alfa = 1.0 / (1.0 + np.exp(-saida))                       # o modelo devolve logits
    mascara = Image.fromarray((alfa * 255).astype(np.uint8)).resize(rgb.size, Image.BILINEAR)
    resultado = rgb.convert("RGBA")
    resultado.putalpha(mascara)
    return resultado


def aumentar(imagem, pedaco=256, borda=16):
    """A foto com 4x a resolução (menos, se passar de LADO_MAXIMO_AUMENTADO), em pedaços para caber na memória."""
    import numpy as np
    from PIL import Image
    rgb = imagem.convert("RGB")
    w, h = rgb.size
    x = np.asarray(rgb, dtype=np.float32) / 255.0
    saida = np.zeros((h * 4, w * 4, 3), dtype=np.float32)
    sessao = _sessao("aumentar")
    for y0 in range(0, h, pedaco):
        for x0 in range(0, w, pedaco):
            ya, xa = max(0, y0 - borda), max(0, x0 - borda)
            yb, xb = min(h, y0 + pedaco + borda), min(w, x0 + pedaco + borda)
            r = sessao.run(None, {"input": x[ya:yb, xa:xb].transpose(2, 0, 1)[None]})[0][0].transpose(1, 2, 0)
            dy, dx = (y0 - ya) * 4, (x0 - xa) * 4
            ph, pw = min(pedaco, h - y0) * 4, min(pedaco, w - x0) * 4
            saida[y0 * 4:y0 * 4 + ph, x0 * 4:x0 * 4 + pw] = r[dy:dy + ph, dx:dx + pw]
    final = Image.fromarray((np.clip(saida, 0, 1) * 255).round().astype(np.uint8))
    if max(final.size) > LADO_MAXIMO_AUMENTADO:
        final.thumbnail((LADO_MAXIMO_AUMENTADO, LADO_MAXIMO_AUMENTADO), Image.LANCZOS)
    return final


def selecionar(imagem, ponto):
    """A máscara (PIL "L", 255 = o objeto) do que está no ponto (x, y) clicado, na resolução da foto."""
    import numpy as np
    from PIL import Image
    rgb = imagem.convert("RGB")
    w, h = rgb.size
    escala = 1024 / max(w, h)
    reduzida = rgb.resize((max(1, round(w * escala)), max(1, round(h * escala))), Image.BILINEAR)
    emb = _sessao("selecionar_codificador").run(None, {"input_image": np.asarray(reduzida, dtype=np.float32)})[0]
    pontos = np.array([[[ponto[0] * escala, ponto[1] * escala], [0, 0]]], dtype=np.float32)
    rotulos = np.array([[1, -1]], dtype=np.float32)          # o ponto clicado e o "ponto vazio" que o SAM pede
    mascaras = _sessao("selecionar_decodificador").run(None, {
        "image_embeddings": emb, "point_coords": pontos, "point_labels": rotulos,
        "mask_input": np.zeros((1, 1, 256, 256), np.float32), "has_mask_input": np.zeros(1, np.float32),
        "orig_im_size": np.array([h, w], np.float32)})[0]
    return Image.fromarray(((mascaras[0, 0] > 0) * 255).astype(np.uint8))


def remover(imagem, mascara, folga=8):
    """A foto sem o que a máscara marca (255 = tirar), preenchido pelo MI-GAN. A máscara cresce `folga` px (a borda do
    objeto não fica de contorno)."""
    import numpy as np
    from PIL import ImageFilter
    rgb = imagem.convert("RGB")
    m = mascara.convert("L").resize(rgb.size).filter(ImageFilter.MaxFilter(folga * 2 + 1))
    manter = (np.asarray(m) < 128).astype(np.uint8) * 255   # no MI-GAN 255 = manter, 0 = buraco (conferido em 25/10)
    from PIL import Image
    saida = _sessao("remover").run(None, {"image": np.asarray(rgb).transpose(2, 0, 1)[None],
                                          "mask": manter[None, None]})[0][0]
    return Image.fromarray(saida.transpose(1, 2, 0))

	"""5.2 (decisão do Enzo, 25/10): os modelinhos de foto, todos baixados obrigatoriamente e rodando na CPU (ONNX Runtime):
	tirar o fundo (BiRefNet lite), aumentar a resolução (Real-ESRGAN general x4v3), selecionar por clique (MobileSAM) e
	remover um objeto (MI-GAN, com a máscara do selecionar). Cada arquivo tem endereço com o commit fixo e o SHA-256
	conferido; o que não confere é apagado. Sem Qt.

	Medido em 25/10 no PC do Enzo (16 threads, só CPU): tirar o fundo ~9 s por foto (3 s de carga na primeira), selecionar
	~0,5 s, remover ~1 s, aumentar ~0,5 s a cada pedaço de 256 px."""
	import hashlib
	import os
	import threading
	import urllib.request
	from pathlib import Path

	import tars_plataforma as plataforma

	_HF = "https://huggingface.co/"
	# nome: (arquivo, endereço com o commit fixo, SHA-256, bytes, licença)
	CATALOGO = {
	"fundo": ("birefnet_lite.onnx",
	_HF + "onnx-community/BiRefNet_lite-ONNX/resolve/de15b22ba131738a16dff04aab8bdf8dc32e3ac1/onnx/model.onnx",
	"5600024376f572a557870a5eb0afb1e5961636bef4e1e22132025467d0f03333", 224005088, "MIT"),
	"aumentar": ("realesr_general_x4v3.onnx",
	_HF + "vlad-grig/realesr-general-x4v3-onnx/resolve/0e95d93e454cf6b5de9bdf5a1c9e7ec340d38231/onnx/model.onnx",
	"7b10e80021ec3e0ff35017ae5022791879d93a799b2e201e1010a37510202266", 4866420, "BSD-3-Clause"),
	"remover": ("migan_pipeline_v2.onnx",
	_HF + "andraniksargsyan/migan/resolve/406830d0fa60666da0071c342ad2fbc8f30c5c64/migan_pipeline_v2.onnx",
	"6f1f3530a1a2324b19752018ce756088b07973cda8d7d890034ace5c8a48c40b", 28079181, "MIT"),
	"selecionar_codificador": ("mobilesam_encoder.onnx",
	_HF + "Acly/MobileSAM/resolve/0d3b403339b4674a82493d5e97964dd78089ddc8/mobile_sam_image_encoder.onnx",
	"580f5fb648ea1062c0aabc26217aed56921985f03f0cbbd852bba81d760cc749", 28157093, "MIT"),
	"selecionar_decodificador": ("mobilesam_decoder.onnx",
	_HF + "Acly/MobileSAM/resolve/0d3b403339b4674a82493d5e97964dd78089ddc8/sam_mask_decoder_single.onnx",
	"93915fc7c993ab9d59ab8c9ccd3bce37f7509c81ab4150a74abd4d2abbd8570d", 16501323, "MIT"),
	}
	LADO_MAXIMO_AUMENTADO = 4096 # aumentar nunca passa disso no lado maior (memória e tempo na CPU)


	def pasta():
	return Path(os.environ.get("TARS_PASTA_MODELINHOS") or plataforma.pastas()["dados"] / "modelinhos")


	def caminho(nome):
	return pasta() / CATALOGO[nome][0]


	def faltando():
	"""Os modelinhos que ainda não estão no disco (pelo tamanho; o hash foi conferido ao baixar). Os testes do app
	(TARS_SEM_MODELINHOS=1) não baixam ~300 MB da internet."""
	if os.environ.get("TARS_SEM_MODELINHOS") == "1":
	return []
	return [n for n, (_, _, _, tam, _) in CATALOGO.items() if not (caminho(n).is_file() and caminho(n).stat().st_size == tam)]


	def bytes_faltando():
	return sum(CATALOGO[n][3] for n in faltando())


	def baixar(nome, ao_progredir=None, cancelado=None, abrir=None):
	"""Baixa um modelinho para um .part, confere o SHA-256 e só então põe no lugar. (ok, erro)."""
	arquivo, url, sha, tamanho, _ = CATALOGO[nome]
	destino = pasta() / arquivo
	destino.parent.mkdir(parents=True, exist_ok=True)
	parcial = destino.with_name(arquivo + ".part")
	soma, feito = hashlib.sha256(), 0
	try:
	with (abrir or urllib.request.urlopen)(urllib.request.Request(url, headers={"User-Agent": "openTARS"}), timeout=60) as r, \
	open(parcial, "wb") as f:
	while True:
	if cancelado and cancelado():
	raise InterruptedError("cancelado")
	bloco = r.read(1 << 20)
	if not bloco:
	break
	f.write(bloco)
	soma.update(bloco)
	feito += len(bloco)
	if ao_progredir:
	ao_progredir(feito, tamanho)
	except InterruptedError:
	parcial.unlink(missing_ok=True)
	return False, "cancelado"
	except OSError as e:
	parcial.unlink(missing_ok=True)
	return False, str(e)
	if soma.hexdigest() != sha:
	parcial.unlink(missing_ok=True)
	return False, f"SHA-256 diferente em {arquivo}"
	os.replace(parcial, destino)
	return True, ""


	def baixar_todos(ao_progredir=None, cancelado=None, abrir=None):
	"""Baixa o que falta, um por um. ao_progredir(feito, total) é do conjunto. (ok, erro)."""
	nomes = faltando()
	total, antes = sum(CATALOGO[n][3] for n in nomes), 0
	for n in nomes:
	ok, erro = baixar(n, (lambda f, _t, a=antes: ao_progredir(a + f, total)) if ao_progredir else None, cancelado, abrir)
	if not ok:
	return False, erro
	antes += CATALOGO[n][3]
	return True, ""


	# ------------------------------------------------------------
	# Os modelos (sessões abertas uma vez e guardadas)
	# ------------------------------------------------------------
	_SESSOES, _TRAVA = {}, threading.Lock()


	def _sessao(nome):
	with _TRAVA:
	if nome not in _SESSOES:
	import onnxruntime
	if not caminho(nome).is_file():
	raise FileNotFoundError(f"modelinho {nome} não baixado")
	_SESSOES[nome] = onnxruntime.InferenceSession(str(caminho(nome)), providers=["CPUExecutionProvider"])
	return _SESSOES[nome]


	def tirar_fundo(imagem):
	"""A foto (PIL) com o fundo transparente (RGBA)."""
	import numpy as np
	from PIL import Image
	rgb = imagem.convert("RGB")
	x = np.asarray(rgb.resize((1024, 1024), Image.BILINEAR), dtype=np.float32) / 255.0
	x = (x - (0.485, 0.456, 0.406)) / (0.229, 0.224, 0.225)
	saida = _sessao("fundo").run(None, {"input_image": x.transpose(2, 0, 1)[None].astype(np.float32)})[0][0, 0]
	alfa = 1.0 / (1.0 + np.exp(-saida)) # o modelo devolve logits
	mascara = Image.fromarray((alfa * 255).astype(np.uint8)).resize(rgb.size, Image.BILINEAR)
	resultado = rgb.convert("RGBA")
	resultado.putalpha(mascara)
	return resultado


	def aumentar(imagem, pedaco=256, borda=16):
	"""A foto com 4x a resolução (menos, se passar de LADO_MAXIMO_AUMENTADO), em pedaços para caber na memória."""
	import numpy as np
	from PIL import Image
	rgb = imagem.convert("RGB")
	w, h = rgb.size
	x = np.asarray(rgb, dtype=np.float32) / 255.0
	saida = np.zeros((h * 4, w * 4, 3), dtype=np.float32)
	sessao = _sessao("aumentar")
	for y0 in range(0, h, pedaco):
	for x0 in range(0, w, pedaco):
	ya, xa = max(0, y0 - borda), max(0, x0 - borda)
	yb, xb = min(h, y0 + pedaco + borda), min(w, x0 + pedaco + borda)
	r = sessao.run(None, {"input": x[ya:yb, xa:xb].transpose(2, 0, 1)[None]})[0][0].transpose(1, 2, 0)
	dy, dx = (y0 - ya) * 4, (x0 - xa) * 4
	ph, pw = min(pedaco, h - y0) * 4, min(pedaco, w - x0) * 4
	saida[y0 * 4:y0 * 4 + ph, x0 * 4:x0 * 4 + pw] = r[dy:dy + ph, dx:dx + pw]
	final = Image.fromarray((np.clip(saida, 0, 1) * 255).round().astype(np.uint8))
	if max(final.size) > LADO_MAXIMO_AUMENTADO:
	final.thumbnail((LADO_MAXIMO_AUMENTADO, LADO_MAXIMO_AUMENTADO), Image.LANCZOS)
	return final


	def selecionar(imagem, ponto):
	"""A máscara (PIL "L", 255 = o objeto) do que está no ponto (x, y) clicado, na resolução da foto."""
	import numpy as np
	from PIL import Image
	rgb = imagem.convert("RGB")
	w, h = rgb.size
	escala = 1024 / max(w, h)
	reduzida = rgb.resize((max(1, round(w * escala)), max(1, round(h * escala))), Image.BILINEAR)
	emb = _sessao("selecionar_codificador").run(None, {"input_image": np.asarray(reduzida, dtype=np.float32)})[0]
	pontos = np.array([[[ponto[0] * escala, ponto[1] * escala], [0, 0]]], dtype=np.float32)
	rotulos = np.array([[1, -1]], dtype=np.float32) # o ponto clicado e o "ponto vazio" que o SAM pede
	mascaras = _sessao("selecionar_decodificador").run(None, {
	"image_embeddings": emb, "point_coords": pontos, "point_labels": rotulos,
	"mask_input": np.zeros((1, 1, 256, 256), np.float32), "has_mask_input": np.zeros(1, np.float32),
	"orig_im_size": np.array([h, w], np.float32)})[0]
	return Image.fromarray(((mascaras[0, 0] > 0) * 255).astype(np.uint8))


	def remover(imagem, mascara, folga=8):
	"""A foto sem o que a máscara marca (255 = tirar), preenchido pelo MI-GAN. A máscara cresce `folga` px (a borda do
	objeto não fica de contorno)."""
	import numpy as np
	from PIL import ImageFilter
	rgb = imagem.convert("RGB")
	m = mascara.convert("L").resize(rgb.size).filter(ImageFilter.MaxFilter(folga * 2 + 1))
	manter = (np.asarray(m) < 128).astype(np.uint8) * 255 # no MI-GAN 255 = manter, 0 = buraco (conferido em 25/10)
	from PIL import Image
	saida = _sessao("remover").run(None, {"image": np.asarray(rgb).transpose(2, 0, 1)[None],
	"mask": manter[None, None]})[0][0]
	return Image.fromarray(saida.transpose(1, 2, 0))