Update README.md

README.md

@@ -44,28 +44,29 @@ This repository contains a ResNet-based convolutional neural network trained to
 import torch
 from torchvision import transforms
 from PIL import Image
+from transformers import pipeline
 
-# Load the model
-model = torch.hub.load('DineshKumar1329/DogCat_Classifier', 'resnet_cat_dog_classifier')
-
-# Define the transformation
+# Define the image transformation
 transform = transforms.Compose([
     transforms.Resize((128, 128)),
     transforms.ToTensor(),
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
 ])
 
-# Load an image
+# Load the model from Hugging Face
+pipe = pipeline("image-classification", model="DineshKumar1329/DogCat_Classifier")
+
+# Load and preprocess an image
 image_path = 'path/to/your/image.jpg'
 image = Image.open(image_path)
 image = transform(image)
 image = image.unsqueeze(0)  # Add batch dimension
 
 # Make a prediction
-model.eval()
-with torch.no_grad():
-    outputs = model(image)
-    temp, predicted = torch.max(outputs, 1)
+result = classifier(image_path)
+
+# Extract the predicted label
+predicted_label = result[0]['label']
 
 # Output the prediction
-print(f'The predicted class for the image is: {"Cat" if predicted.item() == 0 else "Dog"}')
+print(f'The predicted class for the image is: {predicted_label}')