Spaces:
Running
on
Zero
Running
on
Zero
| import gradio as gr | |
| import spaces | |
| from transformers import AutoImageProcessor, SiglipForImageClassification | |
| from transformers.image_utils import load_image | |
| from PIL import Image | |
| import torch | |
| # Load model and processor | |
| model_name = "prithivMLmods/Multisource-121-DomainNet" | |
| model = SiglipForImageClassification.from_pretrained(model_name) | |
| processor = AutoImageProcessor.from_pretrained(model_name) | |
| def multisource_classification(image): | |
| """Predicts the domain category for an input image.""" | |
| # Convert the input numpy array to a PIL Image and ensure it is in RGB format | |
| image = Image.fromarray(image).convert("RGB") | |
| # Process the image and convert it to model inputs | |
| inputs = processor(images=image, return_tensors="pt") | |
| # Get model predictions without gradient calculations | |
| with torch.no_grad(): | |
| outputs = model(**inputs) | |
| logits = outputs.logits | |
| # Convert logits to probabilities using softmax | |
| probs = torch.nn.functional.softmax(logits, dim=1).squeeze().tolist() | |
| # Mapping from class indices to domain labels | |
| labels = { | |
| "0": "barn", "1": "baseball_bat", "2": "basket", "3": "beach", "4": "bear", | |
| "5": "beard", "6": "bee", "7": "bird", "8": "blueberry", "9": "bowtie", | |
| "10": "bracelet", "11": "brain", "12": "bread", "13": "broccoli", "14": "bus", | |
| "15": "butterfly", "16": "circle", "17": "cloud", "18": "cruise_ship", "19": "dolphin", | |
| "20": "dumbbell", "21": "elephant", "22": "eye", "23": "eyeglasses", "24": "feather", | |
| "25": "fish", "26": "flower", "27": "foot", "28": "frog", "29": "giraffe", | |
| "30": "goatee", "31": "golf_club", "32": "grapes", "33": "grass", "34": "guitar", | |
| "35": "hamburger", "36": "hand", "37": "hat", "38": "headphones", "39": "helicopter", | |
| "40": "hexagon", "41": "hockey_stick", "42": "horse", "43": "hourglass", "44": "house", | |
| "45": "ice_cream", "46": "jacket", "47": "ladder", "48": "leg", "49": "lipstick", | |
| "50": "megaphone", "51": "monkey", "52": "moon", "53": "mushroom", "54": "necklace", | |
| "55": "owl", "56": "panda", "57": "pear", "58": "peas", "59": "penguin", | |
| "60": "pig", "61": "pillow", "62": "pineapple", "63": "pizza", "64": "pool", | |
| "65": "popsicle", "66": "rabbit", "67": "rhinoceros", "68": "rifle", "69": "river", | |
| "70": "sailboat", "71": "sandwich", "72": "sea_turtle", "73": "shark", "74": "shoe", | |
| "75": "skyscraper", "76": "snorkel", "77": "snowman", "78": "soccer_ball", "79": "speedboat", | |
| "80": "spider", "81": "spoon", "82": "square", "83": "squirrel", "84": "stethoscope", | |
| "85": "strawberry", "86": "streetlight", "87": "submarine", "88": "suitcase", "89": "sun", | |
| "90": "sweater", "91": "sword", "92": "table", "93": "teapot", "94": "teddy-bear", | |
| "95": "telephone", "96": "tent", "97": "The_Eiffel_Tower", "98": "The_Great_Wall_of_China", | |
| "99": "The_Mona_Lisa", "100": "tiger", "101": "toaster", "102": "tooth", "103": "tornado", | |
| "104": "tractor", "105": "train", "106": "tree", "107": "triangle", "108": "trombone", | |
| "109": "truck", "110": "trumpet", "111": "umbrella", "112": "vase", "113": "violin", | |
| "114": "watermelon", "115": "whale", "116": "windmill", "117": "wine_glass", "118": "yoga", | |
| "119": "zebra", "120": "zigzag" | |
| } | |
| # Create a dictionary mapping each label to its corresponding probability (rounded) | |
| predictions = {labels[str(i)]: round(probs[i], 3) for i in range(len(probs))} | |
| return predictions | |
| # Create Gradio interface | |
| iface = gr.Interface( | |
| fn=multisource_classification, | |
| inputs=gr.Image(type="numpy"), | |
| outputs=gr.Label(label="Prediction Scores"), | |
| title="Multisource-121-DomainNet Classification", | |
| description="Upload an image to classify it into one of 121 domain categories." | |
| ) | |
| # Launch the app | |
| if __name__ == "__main__": | |
| iface.launch() |