app.py

import os
import logging
import gradio as gr
from huggingface_hub import InferenceClient

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger(__name__)

# Environment variables for configuration
HF_TOKEN = os.environ.get("HF_TOKEN", "")
MODEL_ID = os.environ.get("MODEL_ID", "openai/clip-vit-base-patch32")
DEFAULT_LABELS = os.environ.get("DEFAULT_LABELS", "cat, dog, bird, car, person")

logger.info(f"HF_TOKEN configured: {bool(HF_TOKEN)}")
logger.info(f"MODEL_ID: {MODEL_ID}")

client = InferenceClient(token=HF_TOKEN) if HF_TOKEN else InferenceClient()
logger.info("InferenceClient initialized")



def classify(image, labels_text: str) -> dict:
    """Classify image against user-provided labels."""
    logger.info(f"classify() called | image={image is not None} | labels={labels_text[:30] if labels_text else 'None'}")
    
    if image is None:
        logger.warning("No image provided")
        return {}
    
    labels = [l.strip() for l in labels_text.split(",") if l.strip()]
    if not labels:
        labels = [l.strip() for l in DEFAULT_LABELS.split(",")]
    logger.info(f"Labels: {labels}")
    
    try:
        logger.info(f"Calling zero_shot_image_classification | model={MODEL_ID}")
        results = client.zero_shot_image_classification(image, candidate_labels=labels, model=MODEL_ID)
        output = {r.label: r.score for r in results}
        logger.info(f"Results: {output}")
        return output
    except Exception as e:
        logger.error(f"API error: {e}")
        return {"Error": str(e)}


logger.info("Building Gradio interface...")

with gr.Blocks(title="Image Detective") as demo:
    gr.Markdown("# 🔍 Image Detective\nUpload an image and define your own categories!")

    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            img_input = gr.Image(type="pil", label="Upload or drop an image")
            labels_input = gr.Textbox(
                label="Categories (comma-separated)",
                placeholder="cat, dog, bird, car",
                value=DEFAULT_LABELS,
            )
            btn = gr.Button("Classify!", variant="primary")
        with gr.Column(scale=1):
            output = gr.Label(label="Predictions", num_top_classes=5)

    btn.click(classify, inputs=[img_input, labels_input], outputs=output)
    labels_input.submit(classify, inputs=[img_input, labels_input], outputs=output)

demo.queue()
logger.info("Starting Gradio server...")
demo.launch()