Add application file

app.py

@@ -6,7 +6,7 @@ import cv2
 import torch
 import torch.nn as nn
 from PIL import Image
-from torchvision import transforms
+import torchvision
 from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
 from timm.data import create_transform
 
@@ -20,14 +20,14 @@ def pil_loader(filepath):
         img = img.convert('RGB')
     return img
 
-def build_transforms(input_size):
-    transform = transforms.Compose([
-        transforms.Resize(input_size * 8 // 7),
-        transforms.CenterCrop(input_size),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+def build_transforms(input_size, center_crop=True):
+    transform = torchvision.transforms.Compose([
+        torchvision.transforms.Resize(input_size * 8 // 7),
+        torchvision.transforms.CenterCrop(input_size),
+        torchvision.transforms.ToTensor(),
+        torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
              ])
-    return transforms
+    return transform
 
 # Download human-readable labels for Bamboo.
 with open('./trainid2name.json') as f:
@@ -40,11 +40,6 @@ build model
 model = timmvit(pretrain_path='./Bamboo_v0-1_ViT-B16.pth.tar.convert')
 model.eval()
 
-'''
-build data transform
-'''
-eval_transforms = build_transforms(224)
-
 '''
 borrow code from here: https://github.com/jacobgil/pytorch-grad-cam/blob/master/pytorch_grad_cam/utils/image.py
 '''
@@ -73,15 +68,19 @@ def show_cam_on_image(img: np.ndarray,
     # cam = cam / np.max(cam)
     return np.uint8(255 * cam)
 
+
+
+
 def recognize_image(image):
     img_t = eval_transforms(image) 
-
     # compute output
     output = model(img_t.unsqueeze(0))
     prediction = output.softmax(-1).flatten()
     _,top5_idx = torch.topk(prediction, 5)
     return {id2name[str(i)][0]: float(prediction[i]) for i in top5_idx.tolist()}
 
+eval_transforms = build_transforms(224)
+
 
 image = gr.inputs.Image()
 label = gr.outputs.Label(num_top_classes=5)