app.py

import torch
from torchvision import models, transforms
from PIL import Image
import gradio as gr
from huggingface_hub import hf_hub_download

# === Constants ===
REPO_ID = "mvysotskyi/cat_dog_classifier"
FILENAME = "models/cat_dog_classifier.pth"
CLASS_NAMES = ["cat", "dog"]

# === Download model from Hugging Face Hub ===
model_path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)

# === Load model ===
def load_model():
    model = models.mobilenet_v2(pretrained=False)
    model.classifier[1] = torch.nn.Linear(model.last_channel, 2)
    model.load_state_dict(torch.load(model_path, map_location="cpu"))
    model.eval()
    return model

model = load_model()

# === Preprocessing ===
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225]),
])

# === Inference function ===
def classify_image(image):
    image = image.convert("RGB")
    input_tensor = transform(image).unsqueeze(0)

    with torch.no_grad():
        outputs = model(input_tensor)
        probs = torch.softmax(outputs, dim=1)[0]
        result = {CLASS_NAMES[i]: float(probs[i]) for i in range(2)}
    return result

# === Gradio Interface ===
interface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=2),
    title="Cat vs Dog Classifier",
    description="Upload an image of a cat or dog. The model will predict the class with confidence."
)

if __name__ == "__main__":
    interface.launch()