Spaces:
Sleeping
Sleeping
File size: 1,113 Bytes
b8a2709 | 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 |
import os
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as tt
import albumentations as A
from albumentations.pytorch import ToTensorV2
stats = (0.4862, 0.4561, 0.3941), (0.2202, 0.2142, 0.2160)
model_tsfm = A.Compose([
A.Resize(224, 224),
A.Normalize(*stats),
ToTensorV2()
])
classes = ['Australian terrier', 'Border terrier', 'Samoyed', 'Beagle', 'Shih-Tzu', 'English foxhound', 'Rhodesian ridgeback', 'Dingo', 'Golden retriever', 'Old English sheepdog']
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('-i', '--Image',
help="input image path", required=True)
args = vars(parser.parse_args())
print(args)
img_path = args['Image']
#plt.imshow(get_image(img_path, model_tsfm).permute(1,2,0))
#img_pred = eff_b2(get_image(img_path, model_tsfm).unsqueeze(0).to(device))
#print(img_pred)
#img_class = torch.argmax(img_pred)
#print(img_class)
#print(classes[img_class.item()])
|