distance_smiles / app.py
chouchouvs's picture
Update app.py
2b64a85 verified
import os
import shutil
import gradio as gr
from PIL import Image
import tempfile
from pathlib import Path
MODEL_DIR = os.path.join(os.path.dirname(__file__), "models", "decimer")
os.makedirs(MODEL_DIR, exist_ok=True)
os.environ["DECIMER_CACHE_DIR"] = MODEL_DIR
def ensure_model_in_models():
"""Vérifie si le modèle est bien dans /models/decimer, sinon copie depuis le cache utilisateur."""
# Emplacement typique du cache DECIMER (à adapter si besoin)
home = Path.home()
possible_paths = [
home / ".cache" / "decimer",
home / ".data" / "DECIMER-V2"
]
# Test si déjà dans /models/decimer
if any(Path(MODEL_DIR).glob("*")):
return # déjà présent, rien à faire
# Sinon, on force le téléchargement dans le cache via un appel bidon
from DECIMER import predict_SMILES
img_path = os.path.join(MODEL_DIR, "dummy.png")
with open(img_path, "wb") as f:
f.write(b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06\x00\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\nIDATx\xdacd\xf8\x0f\x00\x01\x05\x01\x02\xa2Z\xb2\x00\x00\x00\x00IEND\xaeB`\x82")
try:
predict_SMILES(img_path)
except Exception:
pass
# On cherche le modèle dans le cache
for p in possible_paths:
if p.exists():
# Recopie tous les fichiers/dossiers du cache dans /models/decimer
for src in p.iterdir():
dest = Path(MODEL_DIR) / src.name
if src.is_dir():
if dest.exists():
shutil.rmtree(dest)
shutil.copytree(src, dest)
else:
shutil.copy2(src, dest)
print(f"Modèles copiés de {p} vers {MODEL_DIR}")
ensure_model_in_models()
from DECIMER import predict_SMILES
def predict_fn(img):
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
img.save(tmp, format="PNG")
path = tmp.name
return predict_SMILES(path)
iface = gr.Interface(
fn=predict_fn,
inputs=gr.Image(type="pil"),
outputs="text"
)
iface.launch(share=True)