Update app.py

app.py

@@ -6,48 +6,44 @@ from PIL import Image
 from pytorch_grad_cam import GradCAM
 from pytorch_grad_cam.utils.image import show_cam_on_image
 import gradio as gr
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+import config
+
+from utils import plot_single_image
 
 from model import YOLOv3
 import config
+from torchvision import transforms
+
+scaled_anchors = (
+            torch.tensor(config.ANCHORS)
+            * torch.tensor(config.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)
+        ).to('cpu')
 
 model = YOLOv3(num_classes=config.NUM_CLASSES)
 model.load_state_dict(torch.load("checkpoint.pth.tar", map_location=torch.device('cpu')), strict=False)
 
-inv_normalize = transforms.Normalize(
-    mean=[-0.50/0.23, -0.50/0.23, -0.50/0.23],
-    std=[1/0.23, 1/0.23, 1/0.23]
+test_transforms_exp = A.Compose(
+    [
+        A.LongestMaxSize(max_size=config.IMAGE_SIZE),
+        A.PadIfNeeded(
+            min_height=config.IMAGE_SIZE, min_width=config.IMAGE_SIZE
+        ),
+        A.Normalize(mean=[0, 0, 0], std=[1, 1, 1], max_pixel_value=255,),
+        ToTensorV2(),
+    ]
 )
-classes = ('plane', 'car', 'bird', 'cat', 'deer',
-           'dog', 'frog', 'horse', 'ship', 'truck')
 
 def inference(input_img,show_gradcam="yes", transparency = 0.5, target_layer_number = -1):
+
     transform = transforms.ToTensor()
     org_img = input_img
     input_img = transform(input_img)
     input_img = input_img
     input_img = input_img.unsqueeze(0)
-    model.eval() 
-    outputs = model(input_img)
-    softmax = torch.nn.Softmax(dim=0)
-    #o = softmax(outputs.flatten())
-    #confidences = {classes[i]: float(o[i]) for i in range(10)}
-    #sorted_confidences = {k: v for k, v in sorted(confidences.items(), key=lambda item: item[1], reverse=True)}
-    #sorted_confidences = dict(list(sorted_confidences.items())[:num_classes])    
-    #_, prediction = torch.max(outputs, 1)
-    #target_layers = [model.convblockL3R1[target_layer_number]]
-    #cam = GradCAM(model=model, target_layers=target_layers, use_cuda=False)
-    #grayscale_cam = cam(input_tensor=input_img, targets=None)
-    #grayscale_cam = grayscale_cam[0, :]
-    #img = input_img.squeeze(0)
-    #img = inv_normalize(img)
-    #rgb_img = np.transpose(img, (1, 2, 0))
-    #rgb_img = rgb_img.numpy()
-    #visualization = None
-    #if (show_gradcam == "yes") :        
-    #    visualization = show_cam_on_image(org_img/255, grayscale_cam, use_rgb=True, image_weight=transparency)
-    #else :
-    #    visualization = show_cam_on_image(org_img/255, grayscale_cam, use_rgb=True, image_weight=1)
-    return input_img, org_img
+    out_fig = plot_single_image(model, input_img, 0.6, 0.5,scaled_anchors)
+    return out_fig
 
 title = "TSAI S13 Assignment: YOLO V3 trained on PASCAL VOC Dataset"
 description = "A simple Gradio interface to infer on Custom ResNet model, and get GradCAM results. Please use images that belong to any of these classes - 'plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'."
@@ -56,7 +52,7 @@ examples = [["cat.jpg","yes", 0.5, -1]
 demo = gr.Interface(
     inference, 
     inputs = [gr.Image(shape=(416, 416), label="Input Image"), gr.Radio(["yes", "no"], label="Show Gradcam"),gr.Slider(0, 1, value = 0.5, label="If yes, Opacity of GradCAM"), gr.Slider(-2, -1, value = -2, step=1, label="If yes, Which Layer?")], 
-    outputs = [ gr.Image(shape=(416, 416), label="Output").style(width=128, height=128), gr.Image(shape=(32, 32), label="GradcamOutput").style(width=128, height=128)],
+    outputs = [gr.Plot(label="Plot")],
     title = title,
     description = description,
     examples = examples,