abdulrafay9's picture
Update app.py
ea22a86 verified
Raw
History Blame Contribute Delete
1.34 kB
import os
import importlib.util
from huggingface_hub import hf_hub_download
import gradio as gr
# --- CONFIG ---
PRIVATE_DATASET_ID = "abdulrafay9/containeralign-private"
TOKEN = os.environ.get("HF_TOKEN")
if not TOKEN:
raise RuntimeError("HF_TOKEN is not set. Add it in Settings → Variables and secrets → Secrets.")
# --- DOWNLOAD FILES ---
core_path = hf_hub_download(
repo_id=PRIVATE_DATASET_ID,
repo_type="dataset",
filename="app_core.py",
token=TOKEN,
)
weights_path = hf_hub_download(
repo_id=PRIVATE_DATASET_ID,
repo_type="dataset",
filename="model.pth",
token=TOKEN,
)
# --- LOAD MODULE ---
spec = importlib.util.spec_from_file_location("app_core", core_path)
app_core = importlib.util.module_from_spec(spec)
spec.loader.exec_module(app_core)
# --- LOAD MODEL ---
model = app_core.load_model(weights_path)
# --- DEFINE PREDICTION FUNCTION ---
def predict(image):
return app_core.predict_alignment(model, image)
# --- GRADIO INTERFACE ---
demo = gr.Interface(
fn=predict,
inputs=gr.Image(type="numpy", label="Upload Image"),
outputs=gr.Textbox(label="Prediction Result"),
title="Container Alignment Detection",
description="Upload an image to check whether containers are Aligned or Not Aligned."
)
# --- RUN APP ---
if __name__ == "__main__":
demo.launch()