Spaces:
Sleeping
Sleeping
File size: 4,637 Bytes
419fc11 36fca7c 419fc11 6333fb2 419fc11 36fca7c 419fc11 36fca7c 419fc11 36fca7c 419fc11 36fca7c 419fc11 36fca7c 419fc11 36fca7c 419fc11 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | from __future__ import annotations
from pathlib import Path
from typing import Any, Dict
import gradio as gr
import torch
from PIL import Image
from huggingface_hub import hf_hub_download
import sys
from pathlib import Path as _Path
ROOT = _Path(__file__).resolve().parent
SRC_DIR = ROOT / "src"
if str(SRC_DIR) not in sys.path:
sys.path.insert(0, str(SRC_DIR))
from infer.pipeline import LoadedModel, load_checkpoint, predict_image
MODEL_REPO_ID = "alphamike/GraphAttributeLearning_Model"
MODEL_FILENAMES = {
"baseline": "baseline/best.pt",
"gnn1": "gnn1/best.pt",
"gnn2": "gnn2/best.pt",
}
class ModelRegistry:
def __init__(self) -> None:
self._cache: Dict[str, LoadedModel] = {}
def get(self, key: str, path: Path, device: torch.device) -> LoadedModel:
cached = self._cache.get(key)
if cached is not None:
return cached
loaded = load_checkpoint(path, device=device)
self._cache[key] = loaded
return loaded
registry = ModelRegistry()
def infer_gradio(
image: Image.Image,
model_type: str,
checkpoint_path: str,
top_k: int,
threshold: float,
) -> Any:
if image is None:
return [], "No image provided."
# Resolve checkpoint: either use the textbox value as filename within the model repo
# or fall back to the default mapping for the selected model_type.
filename = (checkpoint_path or "").strip() or MODEL_FILENAMES.get(model_type, "")
if not filename:
return [], f"Unknown model_type '{model_type}'."
try:
local_ckpt_path = hf_hub_download(
repo_id=MODEL_REPO_ID,
filename=filename,
repo_type="model",
)
except Exception as exc: # noqa: BLE001
return [], f"Error downloading checkpoint from Hub: {exc}"
path = Path(local_ckpt_path)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
try:
key = f"{model_type}:{path}"
loaded = registry.get(key, path, device)
# Save temp image to disk-agnostic path is unnecessary; use pipeline logic directly.
from infer.pipeline import predict_image as predict_from_loaded
from infer.pipeline import LoadedModel as _LM # type: ignore
assert isinstance(loaded, _LM)
# Wrap PIL image call using temporary in-memory path emulation.
# For simplicity, reuse adapter directly.
result = loaded.adapter.predict(
image=image,
label_vocab=loaded.label_vocab,
top_k=top_k,
threshold=threshold,
)
rows = [
{"label": label, "score": round(score, 4), "positive@thr": label in result.positives}
for label, score in zip(result.labels, result.scores)
]
message = f"Positives (>= {threshold:.2f}): {', '.join(result.positives) if result.positives else 'none'}"
return rows, message
except Exception as exc: # noqa: BLE001
return [], f"Error: {exc}"
with gr.Blocks() as demo:
gr.Markdown(
"# Adjective-Aware Chair Attributes\n"
"Select baseline, GNN 1, or GNN 2, upload an image, and view attribute scores.\n\n"
f"Models are loaded from Hugging Face model repo: `{MODEL_REPO_ID}`."
)
with gr.Row():
with gr.Column(scale=1):
image_input = gr.Image(type="pil", label="Chair image")
model_type = gr.Radio(
choices=["baseline", "gnn1", "gnn2"],
value="baseline",
label="Model type",
)
checkpoint = gr.Textbox(
value=MODEL_FILENAMES["baseline"],
label="Checkpoint filename (within Hub repo)",
)
def _sync_ckpt(choice: str) -> str:
return MODEL_FILENAMES.get(choice, MODEL_FILENAMES["baseline"])
model_type.change(_sync_ckpt, inputs=model_type, outputs=checkpoint)
top_k = gr.Slider(1, 20, value=5, step=1, label="Top-K")
threshold = gr.Slider(0.0, 1.0, value=0.5, step=0.01, label="Threshold")
run_btn = gr.Button("Run inference")
with gr.Column(scale=2):
table = gr.Dataframe(
headers=["label", "score", "positive@thr"],
datatype=["str", "number", "bool"],
label="Attribute scores",
)
message = gr.Markdown()
run_btn.click(
infer_gradio,
inputs=[image_input, model_type, checkpoint, top_k, threshold],
outputs=[table, message],
)
if __name__ == "__main__":
demo.launch()
|