Spaces:

macadeliccc
/

large-dog-breed-classifier

Sleeping

working

1178181 7 months ago

No virus

1.36 kB

	import spaces
	import gradio as gr
	from transformers import AutoImageProcessor, AutoModelForImageClassification
	import torch
	from PIL import Image

	# Load the fine-tuned model
	model = AutoModelForImageClassification.from_pretrained("Pavarissy/ConvNextV2-large-DogBreed")

	# Initialize the image processor
	preprocessor = AutoImageProcessor.from_pretrained("Pavarissy/ConvNextV2-large-DogBreed")

	def classify_image(image):
	# Preprocess the image
	inputs = preprocessor(images=image, return_tensors="pt")

	# Model prediction
	with torch.no_grad():
	logits = model(**inputs).logits

	# Convert logits to probabilities
	probs = logits.softmax(dim=-1)

	# Extract top 5 predictions
	top_5_probs, top_5_labels = torch.topk(probs, 5)
	top_5_probs = top_5_probs.squeeze().tolist()
	top_5_labels = top_5_labels.squeeze().tolist()

	# Map labels to their names
	labels = model.config.id2label
	predicted_labels = [labels[label] for label in top_5_labels]

	return dict(zip(predicted_labels, top_5_probs))

	# Create a Gradio interface
	iface = gr.Interface(
	fn=classify_image,
	inputs=gr.Image(type="pil"),
	outputs=gr.Label(num_top_classes=5),
	title="Dog Breed Classifier",
	description="Upload an image of a dog, and the model will predict the breed."
	)

	# Launch the interface
	iface.launch(share=True)