Views
No views yet
1from huggingface_hub import PyTorchModelHubMixin
2import torch
3import torch.nn as nn
4from torchvision import models, transforms
5from PIL import Image
6import requests
7
8class CustomDenseNet(nn.Module, PyTorchModelHubMixin):
9 def __init__(self, class_names):
10 super().__init__()
11 self.densenet = models.densenet121(pretrained=True)
12 num_features = self.densenet.classifier.in_features
13 self.densenet.fc = nn.Linear(num_features, len(class_names))
14
15 def forward(self, x):
16 outputs = self.densenet(x)
17 _, preds = torch.max(outputs, 1)
18 probabilities = torch.nn.functional.softmax(outputs, dim=1).squeeze(0)
19
20 predicted_class = class_names[preds.item()]
21 predicted_probabilities = {class_names[i]: probabilities[i].item() for i in range(len(class_names))}
22
23 return predicted_class, predicted_probabilities
24
25model_id = "shinyice/densenet121-dog-emotions"
26class_names = ['angry', 'happy', 'relaxed', 'sad']
27model = CustomDenseNet(class_names)
28model = model.from_pretrained(model_id)
29
30device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
31model.to(device)1def dog_emotion(model, url_mode=False, input_image=None):
2 img_transforms = transforms.Compose([
3 transforms.Resize(224),
4 transforms.ToTensor(),
5 transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
6 if url_mode:
7 image = Image.open(requests.get(input_image, stream=True).raw).convert('RGB')
8 else:
9 image = Image.open(input_image).convert('RGB')
10 image_tensor = img_transforms(image).unsqueeze(0)
11 image_tensor = image_tensor.to(device)
12
13 model.eval()
14 with torch.no_grad():
15 predicted_class, predicted_probabilities = model(image_tensor)
16
17 return predicted_class, predicted_probabilities, image
18
19url_mode = True
20input_image = ""
21
22emotion,probabilities, image = dog_emotion(model=model, url_mode=url_mode, input_image=input_image)
23print(emotion,probabilities)
24image