Create dog_breed.py

dog_breed.py

@@ -0,0 +1,161 @@
+import gradio as gr
+import spaces
+from transformers import AutoImageProcessor, SiglipForImageClassification
+from PIL import Image
+import torch
+
+# Load model and processor
+model_name = "prithivMLmods/Dog-Breed-120"
+model = SiglipForImageClassification.from_pretrained(model_name)
+processor = AutoImageProcessor.from_pretrained(model_name)
+
+@spaces.GPU
+def dog_breed_classification(image):
+    """Predicts the dog breed for an image."""
+    image = Image.fromarray(image).convert("RGB")
+    inputs = processor(images=image, return_tensors="pt")
+    
+    with torch.no_grad():
+        outputs = model(**inputs)
+        logits = outputs.logits
+        probs = torch.nn.functional.softmax(logits, dim=1).squeeze().tolist()
+    
+    labels = {
+        "0": "affenpinscher",
+        "1": "afghan_hound",
+        "2": "african_hunting_dog",
+        "3": "airedale",
+        "4": "american_staffordshire_terrier",
+        "5": "appenzeller",
+        "6": "australian_terrier",
+        "7": "basenji",
+        "8": "basset",
+        "9": "beagle",
+        "10": "bedlington_terrier",
+        "11": "bernese_mountain_dog",
+        "12": "black-and-tan_coonhound",
+        "13": "blenheim_spaniel",
+        "14": "bloodhound",
+        "15": "bluetick",
+        "16": "border_collie",
+        "17": "border_terrier",
+        "18": "borzoi",
+        "19": "boston_bull",
+        "20": "bouvier_des_flandres",
+        "21": "boxer",
+        "22": "brabancon_griffon",
+        "23": "briard",
+        "24": "brittany_spaniel",
+        "25": "bull_mastiff",
+        "26": "cairn",
+        "27": "cardigan",
+        "28": "chesapeake_bay_retriever",
+        "29": "chihuahua",
+        "30": "chow",
+        "31": "clumber",
+        "32": "cocker_spaniel",
+        "33": "collie",
+        "34": "curly-coated_retriever",
+        "35": "dandie_dinmont",
+        "36": "dhole",
+        "37": "dingo",
+        "38": "doberman",
+        "39": "english_foxhound",
+        "40": "english_setter",
+        "41": "english_springer",
+        "42": "entlebucher",
+        "43": "eskimo_dog",
+        "44": "flat-coated_retriever",
+        "45": "french_bulldog",
+        "46": "german_shepherd",
+        "47": "german_short-haired_pointer",
+        "48": "giant_schnauzer",
+        "49": "golden_retriever",
+        "50": "gordon_setter",
+        "51": "great_dane",
+        "52": "great_pyrenees",
+        "53": "greater_swiss_mountain_dog",
+        "54": "groenendael",
+        "55": "ibizan_hound",
+        "56": "irish_setter",
+        "57": "irish_terrier",
+        "58": "irish_water_spaniel",
+        "59": "irish_wolfhound",
+        "60": "italian_greyhound",
+        "61": "japanese_spaniel",
+        "62": "keeshond",
+        "63": "kelpie",
+        "64": "kerry_blue_terrier",
+        "65": "komondor",
+        "66": "kuvasz",
+        "67": "labrador_retriever",
+        "68": "lakeland_terrier",
+        "69": "leonberg",
+        "70": "lhasa",
+        "71": "malamute",
+        "72": "malinois",
+        "73": "maltese_dog",
+        "74": "mexican_hairless",
+        "75": "miniature_pinscher",
+        "76": "miniature_poodle",
+        "77": "miniature_schnauzer",
+        "78": "newfoundland",
+        "79": "norfolk_terrier",
+        "80": "norwegian_elkhound",
+        "81": "norwich_terrier",
+        "82": "old_english_sheepdog",
+        "83": "otterhound",
+        "84": "papillon",
+        "85": "pekinese",
+        "86": "pembroke",
+        "87": "pomeranian",
+        "88": "pug",
+        "89": "redbone",
+        "90": "rhodesian_ridgeback",
+        "91": "rottweiler",
+        "92": "saint_bernard",
+        "93": "saluki",
+        "94": "samoyed",
+        "95": "schipperke",
+        "96": "scotch_terrier",
+        "97": "scottish_deerhound",
+        "98": "sealyham_terrier",
+        "99": "shetland_sheepdog",
+        "100": "shih-tzu",
+        "101": "siberian_husky",
+        "102": "silky_terrier",
+        "103": "soft-coated_wheaten_terrier",
+        "104": "staffordshire_bullterrier",
+        "105": "standard_poodle",
+        "106": "standard_schnauzer",
+        "107": "sussex_spaniel",
+        "108": "test",
+        "109": "tibetan_mastiff",
+        "110": "tibetan_terrier",
+        "111": "toy_poodle",
+        "112": "toy_terrier",
+        "113": "vizsla",
+        "114": "walker_hound",
+        "115": "weimaraner",
+        "116": "welsh_springer_spaniel",
+        "117": "west_highland_white_terrier",
+        "118": "whippet",
+        "119": "wire-haired_fox_terrier",
+        "120": "yorkshire_terrier"
+    }
+    
+    predictions = {labels[str(i)]: round(probs[i], 3) for i in range(len(probs))}
+    return predictions
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=dog_breed_classification,
+    inputs=gr.Image(type="numpy"),
+    outputs=gr.Label(label="Prediction Scores"),
+    title="Dog Breed Classification",
+    description="Upload an image to classify it into one of the 121 dog breed categories."
+)
+
+# Launch the app
+if __name__ == "__main__":
+    iface.launch()