app.py

import gradio as gr
from PIL import Image
import torch
import numpy as np
from transformers import CLIPProcessor, CLIPModel

# Load model and processor
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch16").to(device)
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch16")

# (You can keep your STANFORD_BREEDS and BREED_LIFESPAN dictionaries as-is)

# -------- Helper Functions from your code --------
# Stanford Dogs Dataset - 120 Breeds (from search results [1][3][4])
STANFORD_BREEDS = [
    "afghan hound", "african hunting dog", "airedale",
    "american staffordshire terrier", "appenzeller", "australian terrier",
    "basenji", "basset", "beagle", "bedlington terrier", "bernese mountain dog",
    "black-and-tan coonhound", "blenheim spaniel", "bloodhound", "bluetick",
    "border collie", "border terrier", "borzoi", "boston bull",
    "bouvier des flandres", "boxer", "brabancon griffon", "briard",
    "brittany spaniel", "bull mastiff", "cairn", "cardigan",
    "chesapeake bay retriever", "chihuahua", "chow", "clumber",
    "cocker spaniel", "collie", "curly-coated retriever", "dandie dinmont",
    "dhole", "dingo", "doberman", "english foxhound", "english setter",
    "english springer", "entlebucher", "eskimo dog", "flat-coated retriever",
    "french bulldog", "german shepherd", "german short-haired pointer",
    "giant schnauzer", "golden retriever", "gordon setter", "great dane",
    "great pyrenees", "greater swiss mountain dog", "groenendael",
    "ibizan hound", "irish setter", "irish terrier", "irish water spaniel",
    "irish wolfhound", "italian greyhound", "japanese spaniel", "keeshond",
    "kelpie", "kerry blue terrier", "komondor", "kuvasz", "labrador retriever",
    "lakeland terrier", "leonberg", "lhasa", "malamute", "malinois",
    "maltese dog", "mexican hairless", "miniature pinscher",
    "miniature poodle", "miniature schnauzer", "newfoundland",
    "norfolk terrier", "norwegian elkhound", "norwich terrier",
    "old english sheepdog", "otterhound", "papillon", "pekinese", "pembroke",
    "pomeranian", "pug", "redbone", "rhodesian ridgeback", "rottweiler",
    "saint bernard", "saluki", "samoyed", "schipperke", "scotch terrier",
    "scottish deerhound", "sealyham terrier", "shetland sheepdog", "shih tzu",
    "siberian husky", "silky terrier", "soft-coated wheaten terrier",
    "staffordshire bullterrier", "standard poodle", "standard schnauzer",
    "sussex spaniel", "tibetan mastiff", "tibetan terrier", "toy poodle",
    "toy terrier", "vizsla", "walker hound", "weimaraner",
    "welsh springer spaniel", "west highland white terrier", "whippet",
    "wire-haired fox terrier", "yorkshire terrier"
]

# Breed-specific lifespan data (years) for age prediction
BREED_LIFESPAN = {
    # Sourced from Nature study (2022) and Scientific Reports (2024)
    "afghan hound": 11.1,
    "african hunting dog": 10.5,  # Estimated (similar to other hounds)
    "airedale": 11.5,
    "american staffordshire terrier": 12.5,
    "appenzeller": 13.0,  # Estimated (similar to Swiss breeds)
    "australian terrier": 13.5,
    "basenji": 12.1,
    "basset": 12.5,
    "beagle": 12.5,
    "bedlington terrier": 13.7,
    "bernese mountain dog": 10.1,
    "black-and-tan coonhound": 10.8,
    "blenheim spaniel": 13.3,  # Cavalier King Charles variant
    "bloodhound": 9.3,
    "bluetick": 11.0,
    "border collie": 13.1,
    "border terrier": 14.2,
    "borzoi": 12.0,
    "boston bull": 11.8,
    "bouvier des flandres": 11.3,
    "boxer": 11.3,
    "brabancon griffon": 13.0,
    "briard": 12.6,
    "brittany spaniel": 13.5,
    "bull mastiff": 10.2,
    "cairn": 14.0,
    "cardigan": 13.2,  # Welsh Corgi
    "chesapeake bay retriever": 11.6,
    "chihuahua": 11.8,
    "chow": 12.1,
    "clumber": 12.3,
    "cocker spaniel": 13.3,
    "collie": 13.3,
    "curly-coated retriever": 12.2,
    "dandie dinmont": 12.8,
    "dhole": 10.0,  # Estimated (wild dog)
    "dingo": 10.0,  # Estimated (wild dog)
    "doberman": 11.2,
    "english foxhound": 13.0,
    "english setter": 13.1,
    "english springer": 13.5,
    "entlebucher": 13.0,
    "eskimo dog": 11.3,
    "flat-coated retriever": 11.7,
    "french bulldog": 9.8,
    "german shepherd": 11.3,
    "german short-haired pointer": 13.4,
    "giant schnauzer": 12.1,
    "golden retriever": 13.2,
    "gordon setter": 12.4,
    "great dane": 10.6,
    "great pyrenees": 10.9,
    "greater swiss mountain dog": 10.9,
    "groenendael": 12.0,  # Belgian Shepherd
    "ibizan hound": 13.3,
    "irish setter": 12.9,
    "irish terrier": 13.5,
    "irish water spaniel": 10.8,
    "irish wolfhound": 9.9,
    "italian greyhound": 14.0,
    "japanese spaniel": 13.3,  # Japanese Chin
    "keeshond": 12.3,
    "kelpie": 12.0,
    "kerry blue terrier": 12.4,
    "komondor": 10.5,
    "kuvasz": 10.5,
    "labrador retriever": 13.1,
    "lakeland terrier": 14.2,
    "leonberg": 10.0,
    "lhasa": 14.0,
    "malamute": 11.3,
    "malinois": 12.0,
    "maltese dog": 13.1,
    "mexican hairless": 13.0,  # Xoloitzcuintli
    "miniature pinscher": 13.7,
    "miniature poodle": 14.0,
    "miniature schnauzer": 13.3,
    "newfoundland": 11.0,
    "norfolk terrier": 13.5,
    "norwegian elkhound": 13.0,
    "norwich terrier": 14.0,
    "old english sheepdog": 12.1,
    "otterhound": 12.0,
    "papillon": 14.5,
    "pekinese": 13.3,
    "pembroke": 13.2,  # Welsh Corgi
    "pomeranian": 12.2,
    "pug": 11.6,
    "redbone": 12.0,
    "rhodesian ridgeback": 12.0,
    "rottweiler": 10.6,
    "saint bernard": 9.3,
    "saluki": 13.3,
    "samoyed": 13.1,
    "schipperke": 14.2,
    "scotch terrier": 12.7,  # Scottish Terrier
    "scottish deerhound": 10.5,
    "sealyham terrier": 13.1,
    "shetland sheepdog": 13.4,
    "shih tzu": 12.8,
    "siberian husky": 11.9,
    "silky terrier": 13.3,
    "soft-coated wheaten terrier": 13.7,
    "staffordshire bullterrier": 12.0,
    "standard poodle": 14.0,
    "standard schnauzer": 13.0,
    "sussex spaniel": 13.5,
    "tibetan mastiff": 13.3,
    "tibetan terrier": 13.8,
    "toy poodle": 14.0,
    "toy terrier": 13.0,
    "vizsla": 13.5,
    "walker hound": 12.0,
    "weimaraner": 12.8,
    "welsh springer spaniel": 14.0,
    "west highland white terrier": 13.4,
    "whippet": 13.4,
    "wire-haired fox terrier": 13.5,
    "yorkshire terrier": 13.3
}

def predict_biological_age(image, breed):
    avg_lifespan = BREED_LIFESPAN.get(breed.lower(), 12)
    age_prompts = [f"a {age}-year-old {breed}" for age in range(1, int(avg_lifespan * 2) + 1)]

    inputs = processor(text=age_prompts, images=image, return_tensors="pt", padding=True).to(device)

    with torch.no_grad():
        outputs = model(**inputs)
        age_logits = outputs.logits_per_image
        age_probs = age_logits.softmax(dim=1).cpu().numpy()[0]

    return np.argmax(age_probs) + 1


def analyze_dog_health(image, user_breed=None, chronological_age=None):
    image = image.convert("RGB")
    inputs = processor(images=image, return_tensors="pt").to(device)

    with torch.no_grad():
        image_features = model.get_image_features(**inputs)

    # 1. Breed classification
    breed_inputs = processor(
        text=[f"a photo of a {b}" for b in STANFORD_BREEDS],
        return_tensors="pt",
        padding=True
    ).to(device)

    with torch.no_grad():
        text_features = model.get_text_features(**breed_inputs)
        breed_similarities = (image_features @ text_features.T).softmax(dim=-1)

    top_breed_idx = breed_similarities.argmax().item()
    predicted_breed = STANFORD_BREEDS[top_breed_idx]
    breed_confidence = breed_similarities[0, top_breed_idx].item()

    breed = user_breed if user_breed else predicted_breed

    # 2. Health Aspects
    health_aspects = {
        "coat_health": ["shiny healthy coat", "dull patchy fur"],
        "eye_clarity": ["bright clear eyes", "cloudy milky eyes"],
        "body_condition": ["ideal muscle tone", "visible ribs or hip bones"],
        "dental_health": ["clean white teeth", "yellow stained teeth"],
        "energy_level": ["alert energetic posture", "lethargic tired appearance"],
        "skin_health": ["smooth healthy skin", "red irritated skin"],
        "weight_status": ["healthy weight", "overweight obese"],
        "joint_health": ["strong stable joints", "stiff painful movement"]
    }

    health_report = {}
    for aspect, (pos, neg) in health_aspects.items():
        aspect_input = processor(text=[pos, neg], return_tensors="pt", padding=True).to(device)
        with torch.no_grad():
            aspect_features = model.get_text_features(**aspect_input)
            similarity = (image_features @ aspect_features.T).softmax(dim=-1)
        health_report[aspect] = {
            "assessment": pos if similarity[0, 0] > similarity[0, 1] else neg,
            "confidence": similarity[0, 0].item(),
            "score": similarity[0, 0].item() - similarity[0, 1].item()
        }

    # 3. Biological age
    bio_age = predict_biological_age(image, breed)
    lifespan = BREED_LIFESPAN.get(breed.lower(), 12)
    pace = round(bio_age / chronological_age, 2) if chronological_age else None

    # 4. Final report
    summary = f"Predicted Breed: {predicted_breed} ({breed_confidence:.1%})\n"
    if user_breed:
        summary += f"User Override Breed: {user_breed}\n"
    summary += f"Biological Age: {bio_age} years\n"
    summary += f"Chronological Age: {chronological_age if chronological_age else 'Not provided'}\n"
    summary += f"Avg Lifespan: {lifespan} years\n"
    if pace:
        summary += f"Pace of Aging: {pace:.2f}x\n"

    summary += "\n### Health Assessment\n"
    for k, v in health_report.items():
        summary += f"- {k.replace('_', ' ').title()}: {v['assessment']} (Confidence: {v['confidence']:.1%}, Score: {v['score']:.2f})\n"

    return summary

# -------- Gradio Interface --------

iface = gr.Interface(
    fn=analyze_dog_health,
    inputs=[
        gr.Image(type="pil", label="Upload Dog Image"),
        gr.Textbox(label="(Optional) Override Breed"),
        gr.Number(label="Chronological Age (in years)", precision=1)
    ],
    outputs=gr.Markdown(),
    title="🐶 Dog Health & Age Analyzer",
    description="Upload a photo of your dog to predict its breed, biological age, health indicators, and compare with average lifespan.",
    allow_flagging="never"
)

iface.launch()