# app.py (for your Hugging Face Space/Model Repo: Hajorda/keduClassifier) import gradio as gr import torch import pytorch_lightning as pl from timm import create_model import torch.nn as nn from box import Box import albumentations as A from albumentations.pytorch.transforms import ToTensorV2 import cv2 import pickle from PIL import Image import numpy as np import os # import requests # Commenting out as Giphy API key is not used by default import random # For random choice of keywords if you enable Giphy later from huggingface_hub import hf_hub_download # --- Model and Repository Configuration --- # This should exactly match your model repository on Hugging Face HF_USERNAME = "Hajorda" HF_MODEL_NAME = "keduClassifier" # CORRECTED: Matches your repo name REPO_ID = f"{HF_USERNAME}/{HF_MODEL_NAME}" # --- Inference Configuration --- cfg_dict_for_inference = { 'model_name': 'swin_tiny_patch4_window7_224', # Should match your trained model 'dropout_backbone': 0.1, # Should match your trained model 'dropout_fc': 0.2, # Should match your trained model 'img_size': (224, 224), 'num_classes': 37, # This MUST match the number of classes your model was trained on } cfg_inference = Box(cfg_dict_for_inference) # --- PyTorch Lightning Model Definition --- class PetBreedModel(pl.LightningModule): def __init__(self, cfg: Box): super().__init__() self.cfg = cfg self.backbone = create_model( self.cfg.model_name, pretrained=False, num_classes=0, in_chans=3, drop_rate=self.cfg.dropout_backbone ) # Ensure img_size is a tuple for unpacking h, w = self.cfg.img_size if isinstance(self.cfg.img_size, tuple) else (224, 224) dummy_input = torch.randn(1, 3, h, w) with torch.no_grad(): num_features = self.backbone(dummy_input).shape[-1] self.fc = nn.Sequential( nn.Linear(num_features, num_features // 2), nn.ReLU(), nn.Dropout(self.cfg.dropout_fc), nn.Linear(num_features // 2, self.cfg.num_classes) ) def forward(self, x): features = self.backbone(x) output = self.fc(features) return output # --- Helper Functions to Load Assets from Hugging Face Hub --- def load_model_from_hf_for_space(repo_id=REPO_ID, ckpt_filename="pytorch_model.ckpt"): model_path = hf_hub_download(repo_id=repo_id, filename=ckpt_filename) if cfg_inference.num_classes is None: # Should be set by cfg_dict_for_inference raise ValueError("num_classes must be set in cfg_inference to load the model for Gradio.") # Pass the cfg for the model structure loaded_model = PetBreedModel.load_from_checkpoint(model_path, cfg=cfg_inference, strict=False) loaded_model.eval() return loaded_model def load_label_encoder_from_hf_for_space(repo_id=REPO_ID, le_filename="label_encoder.pkl"): le_path = hf_hub_download(repo_id=repo_id, filename=le_filename) with open(le_path, 'rb') as f: label_encoder = pickle.load(f) return label_encoder # --- Load Model and Label Encoder (once at app startup) --- print(f"Loading model and label encoder from repository: {REPO_ID}") try: model = load_model_from_hf_for_space() label_encoder = load_label_encoder_from_hf_for_space() device = "cuda" if torch.cuda.is_available() else "cpu" model.to(device) print(f"Model and label encoder loaded successfully. Using device: {device}") except Exception as e: print(f"Error loading model or label encoder: {e}") # If loading fails, the Gradio app might not work. # Consider how to handle this, e.g., display an error in the UI. model = None label_encoder = None device = "cpu" # --- Funny GIF Logic --- # funny_cat_keywords = ["funny cat", "silly cat", "cat meme", "derp cat"] # GIPHY_API_KEY = "YOUR_GIPHY_API_KEY" # Optional def get_funny_cat_gif(breed_name): # Using a predefined list for simplicity and to avoid API key requirements predefined_gifs = { "abyssinian": "https://media.giphy.com/media/v1.Y2lkPTc5MGI3NjExaWN4bDNzNWVzM2VqNHE4Ym5zN2ZzZHF0Zzh0bGRqZzRjMnhsZW5pZCZlcD12MV9pbnRlcm5hbF9naWZfYnlfaWQmY3Q9Zw/3oriO0OEd9QIDdllqo/giphy.gif", "american bulldog": "https://media.giphy.com/media/v1.Y2lkPTc5MGI3NjExbHgzYXB6N3g5NThnaXU2eWR2aHljOXg3NjMzbGJwNmF6NmxkdXU2ayZlcD12MV9pbnRlcm5hbF9naWZfYnlfaWQmY3Q9Zw/1simplexLKhMTqI/giphy.gif", # Example for a dog breed "bengal": "https://media.giphy.com/media/v1.Y2lkPTc5MGI3NjExbnl0Z2J6cWtub29qdjFlajQ4ZXZ6czY2ZDY0cW53b3I2amI0OHhoYSZlcD12MV9pbnRlcm5hbF9naWZfYnlfaWQmY3Q9Zw/BK1 SANT0sqq1q/giphy.gif", "birman": "https://media.giphy.com/media/v1.Y2lkPTc5MGI3NjExZ3Q4NXZmMjQ1azE2dHZ2czZnNnBoNThkZ3FkY2Z0c3hqNjVqMTdhaSZlcD12MV9pbnRlcm5hbF9naWZfYnlfaWQmY3Q9Zw/catdogcessing/giphy.gif", "bombay": "https://media.giphy.com/media/v1.Y2lkPTc5MGI3NjExc3N5b2c3MmgwN3JzbjRkYmdocjdhcDc3ejExZGZqZmZtbDBxdXRrcSZlcD12MV9pbnRlcm5hbF9naWZfYnlfaWQmY3Q9Zw/q1MeAPDDMb43K/giphy.gif", "british shorthair": "https://media.giphy.com/media/v1.Y2lkPTc5MGI3NjExYTY3NG96bTc0bnFyOGNkaXBwcTYwdGZzZ3JwY2pscGNmbmZydG05eSZlcD12MV9pbnRlcm5hbF9naWZfYnlfaWQmY3Q9Zw/Lq0h93752f6J9tij39/giphy.gif", "egyptian mau": "https://media.giphy.com/media/v1.Y2lkPTc5MGI3NjExbjZ6dmJvaDhsb3N4ZXdkOXNrbzRkYnJmMHo3MnE2bWJocjU0Mm5jayZlcD12MV9pbnRlcm5hbF9naWZfYnlfaWQmY3Q9Zw/3o7ZeLambpFh3TS2ZO/giphy.gif", "maine coon": "https://media.giphy.com/media/v1.Y2lkPTc5MGI3NjExd3F6NWoyanFmY2xmcHBtMHRhMXAzaXZrYnJia3UxcDRtcXFsYjE2NSZlcD12MV9pbnRlcm5hbF_naWZfYnlfaWQmY3Q9Zw/MDrmyLuUh9A1a/giphy.gif", "persian": "https://media.giphy.com/media/v1.Y2lkPTc5MGI3NjExYW12cDRuc3ZtZ2ZpN2Q2cjdwMHBmb2F3MzJ5d295dGRscG9hdmFpNiZlcD12MV9pbnRlcm5hbF9naWZfYnlfaWQmY3Q9Zw/uE4gVmbjaZmmY/giphy.gif", "ragdoll": "https://media.giphy.com/media/v1.Y2lkPTc5MGI3NjExczZqNWs2ZWU1ZTVobXVxdTZrN2hzcGZoaDVrYnNpZGF4a3FpM3N4aCZlcD12MV9pbnRlcm5hbF9naWZfYnlfaWQmY3Q9Zw/ObTT5h01Xo43C/giphy.gif", "russian blue": "https://media.giphy.com/media/v1.Y2lkPTc5MGI3NjExc3NqcHgzcnVldjA2MnQxc3oyZnp5a2R1eXZxY21hZTN4NHAwd2NyNyZlcD12MV9pbnRlcm5hbF_naWZfYnlfaWQmY3Q9Zw/114ZzmjHizvdsY/giphy.gif", "siamese": "https://media.giphy.com/media/v1.Y2lkPTc5MGI3NjExa3g0dHZtZmRncWN0cnZkNnVnMGRtYjN2ajZ2d3o1cHZtaW50ZHQ5ayZlcD12MV9pbnRlcm5hbF9naWZfYnlfaWQmY3Q9Zw/ICOgUNjpvO0PC/giphy.gif", "sphynx": "https://media.giphy.com/media/v1.Y2lkPTc5MGI3NjExcXZjdzFybXh0ZW53OHI4ZWQxazNtb3N4dDNzOGJrdmZrdXFzbnUyZSZlcD12MV9pbnRlcm5hbF9naWZfYnlfaWQmY3Q9Zw/mlvseq9yvZhba/giphy.gif", "default": "https://media.giphy.com/media/v1.Y2lkPTc5MGI3NjExNWMwNnU4NW9nZTV5c3Z0eThsOHhsOWN0Nnh0a3VzbjFxeGU0bjFuNiZlcD12MV9pbnRlcm5hbF9naWZfYnlfaWQmY3Q9Zw/BzyTuYCmvSORqs1ABM/giphy.gif" } # Normalize breed name for lookup normalized_breed_name = breed_name.lower().replace(" ", "_").replace("-", "_") return predefined_gifs.get(normalized_breed_name, predefined_gifs["default"]) # --- Gradio Interface Function --- def classify_cat_breed(image_input_bgr): # Gradio image is usually BGR numpy array if model is None or label_encoder is None: return ("Model not loaded. Please check logs.", "Error: Model components failed to load.", "") # Convert BGR to RGB img_rgb = cv2.cvtColor(image_input_bgr, cv2.COLOR_BGR2RGB) h, w = cfg_inference.img_size transforms_gradio = A.Compose([ A.Resize(height=h, width=w), A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ToTensorV2(), ]) input_tensor = transforms_gradio(image=img_rgb)['image'].unsqueeze(0).to(device) with torch.no_grad(): logits = model(input_tensor) probabilities = torch.softmax(logits, dim=1) confidence, predicted_idx = torch.max(probabilities, dim=1) predicted_breed_id = predicted_idx.item() predicted_breed_name = label_encoder.inverse_transform([predicted_breed_id])[0] conf_score = confidence.item() funny_message = f"My AI brain (all {conf_score*100:.1f}% of it that's sure) says this purrfect creature is a **{predicted_breed_name}**!" if conf_score < 0.5: funny_message += " ...Though, to be honest, it could also be a very fluffy potato. My circuits are confused! 🥔" elif conf_score < 0.8: funny_message += " Pretty confident, but if it starts barking, don't blame me! 😜" else: funny_message += " Absolutely magnificent! A textbook example, if cats read textbooks. 🧐" gif_url = get_funny_cat_gif(predicted_breed_name) return ( f"**{predicted_breed_name.title()}** (Confidence: {conf_score*100:.2f}%)", funny_message, gif_url ) # --- Define the Gradio Interface --- title = "😼 KEDU's Kompletely Kooky Kat (and K9?) Klassifier! 🐶" description = ( "Upload a pic of your furry overlord (cat OR dog from the Oxford-IIIT set!), and I'll " "attempt a hilariously 'accurate' breed guess. Powered by Swin Transformers and an " "unhealthy obsession with pets. Results may vary, giggles guaranteed!" ) # Corrected article link article_link_href = f"https://huggingface.co/{REPO_ID}" # Uses the correctly defined REPO_ID article = f"

Model based on Swin Transformer, fine-tuned on the Oxford-IIIT Pet Dataset. Model Card & Files

" # Add some example images to your repo and reference them here # For example, if you add 'cat_example.jpg' and 'dog_example.jpg' to your HF repo example_images = [ ["cat1.webp"], # You'll need to upload this image to your HF repo ["cat2.webp"], # You'll need to upload this image to your HF repo ["cat3.webp"], ] # Check if example files exist, if not, provide placeholders or skip examples # This check would ideally be done by trying to download them if they are remote URLs # For local paths in a repo, Gradio handles it if the files are present. iface = gr.Interface( fn=classify_cat_breed, inputs=gr.Image(type="numpy", label="Upload Your Pet's Most Glamorous Shot! 📸"), outputs=[ gr.Textbox(label="🧐 The AI's Verdict Is... (Breed & Confidence)"), gr.Markdown(label="💬 AI's Deep (and Silly) Thoughts..."), # Markdown for bolding gr.Image(type="filepath", label="🎉 Celebration/Confusion GIF! 🎉") ], title=title, description=description, article=article, # examples=example_images, # Uncomment if you add example images to your repo theme=gr.themes.Monochrome(), # Trying a different theme allow_flagging='never' ) if __name__ == "__main__": # When running locally (e.g., python app.py), this will launch the server. # On Hugging Face Spaces, Spaces handles the launch. iface.launch()