-
Notifications
You must be signed in to change notification settings - Fork 1
/
inference.py
53 lines (37 loc) · 1.46 KB
/
inference.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
import json
import torch
from commons import get_model, get_tensor
import os
from torchvision import datasets
import torchvision.transforms as transforms
train_dir = 'images/train'
train_transforms = transforms.Compose([transforms.RandomRotation(10),
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])])
#train_data = datasets.ImageFolder(train_dir , transform=train_transforms)
'''with open('cat_to_name.json') as f:
cat_to_name = json.load(f)
with open('class_to_idx.json') as f:
class_to_idx = json.load(f)
idx_to_class = {v:k for k, v in class_to_idx.items()}'''
model = get_model()
class_names =['kurti', 'saree', 'shirt']
#class_names = [item[4:].replace("_", " ") for item in train_data.classes]
#fp=train_data.classes
def get_flower_name(image_bytes):
tensor = get_tensor(image_bytes)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.eval()
with torch.no_grad():
out = model(tensor.to(device))
ps = torch.exp(out)
top_p, top_class = ps.topk(1, dim=1)
index = top_class.item()
return class_names[index]
'''outputs = model.forward(tensor)
_, prediction = outputs.max(1)
category = prediction.item()
class_idx = idx_to_class[category]
flower_name = cat_to_name[class_idx]
return category, flower_name'''