import numpy as np
from PIL import Image
from tkinter import Tk, filedialog
from transformers import AutoModelForImageClassification, AutoFeatureExtractor

# Load the model and feature extractor from Hugging Face
MODEL_NAME = "shaktibiplab/Animal-Classification"
model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
extractor = AutoFeatureExtractor.from_pretrained(MODEL_NAME)

# Function to load and preprocess the image
def load_and_preprocess_image(image_path):
    image = Image.open(image_path).convert("RGB")
    return image

# Function to classify the image
def classify_image(image_path):
    image = load_and_preprocess_image(image_path)
    inputs = extractor(images=image, return_tensors="pt")
    outputs = model(**inputs)
    logits = outputs.logits
    predicted_class_idx = logits.argmax(-1).item()
    return model.config.id2label[predicted_class_idx]

# Main program with file upload dialog
if __name__ == "__main__":
    root = Tk()
    root.withdraw()  # Hide the main tkinter window
    print("Please select an image file.")
    
    # Open a file dialog to select the image
    image_path = filedialog.askopenfilename(
        title="Select an Image",
        filetypes=[("Image Files", "*.jpg *.jpeg *.png *.bmp *.tiff")]
    )
    
    if image_path:
        try:
            predicted_class = classify_image(image_path)
            print(f"Predicted Class: {predicted_class}")
        except Exception as e:
            print(f"Error: {e}")
    else:
        print("No file selected.")