Update app.py

app.py

@@ -40,31 +40,23 @@ def resize_image_pil(image, new_width, new_height):
 
     return resized
 
-def inference(input_image, transparency=0.5, target_layer_number=-1):
-    input_image = resize_image_pil(input_image, 32, 32)
-    input_image = np.array(input_image)
-    org_img = input_image
-    input_image = input_image.reshape((32, 32, 3))
+def inference(input_img, transparency):
     transform = transforms.ToTensor()
-    input_image = transform(input_image)
-    input_image = input_image.unsqueeze(0)
-    outputs = model(input_image)
-    softmax = torch.nn.Softmax(dim=0)
-    o = softmax(outputs.flatten())
-    confidences = {classes[i]: float(o[i]) for i in range(10)}
+    input_img = transform(input_img)
+    input_img = input_img.to(device)
+    input_img = input_img.unsqueeze(0)
+    outputs = model(input_img)
     _, prediction = torch.max(outputs, 1)
-    target_layers = [model.layer2[target_layer_number]]
-    cam = GradCAM(model= model, target_layers = target_layers)
-    grayscale_cam = cam(input_tensor=input_image, targets=None)
+    target_layers = [model.layer2[-2]]
+    cam = GradCAM(model=model, target_layers=target_layers, use_cuda=True)
+    grayscale_cam = cam(input_tensor=input_img, targets=targets)
     grayscale_cam = grayscale_cam[0, :]
-    visualization = show_cam_on_image(
-        org_img/255, 
-        grayscale_cam, 
-        use_rgb=True,
-        image_weight = transparency
-    )
-     
-    return classes[prediction[0].item(), visualization, confidences]
+    img = input_img.squeeze(0).to('cpu')
+    img = inv_normalize(img)
+    rgb_img = np.transpose(img, (1, 2, 0))
+    rgb_img = rgb_img.numpy()
+    visualization = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True, image_weight=transparency)
+    return classes[prediction[0].item()], visualization
 
 demo = gr.Interface(
     inference,