Update app.py

app.py

@@ -21,7 +21,7 @@ model = FocalNet(depths=[12], patch_size=16, embed_dim=768, focal_levels=[3], us
 url = 'https://projects4jw.blob.core.windows.net/focalnet/release/classification/focalnet_base_iso_16.pth'
 checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
 model.load_state_dict(checkpoint["model"])
-model = model.cuda(); model.eval()
+model.eval()
 
 '''
 build data transform
@@ -65,32 +65,32 @@ def show_cam_on_image(img: np.ndarray,
 def classify_image(inp):
     
     img_t = eval_transforms(inp) 
-    img_d = display_transforms(inp).permute(1, 2, 0).cpu().numpy()
+    img_d = display_transforms(inp).permute(1, 2, 0).numpy()
     print(img_d.min(), img_d.max())
 
-    prediction = model(img_t.unsqueeze(0).cuda()).softmax(-1).flatten()
+    prediction = model(img_t.unsqueeze(0)).softmax(-1).flatten()
 
     modulator = model.layers[0].blocks[2].modulation.modulator.norm(2, 1, keepdim=True)
     modulator = nn.Upsample(size=img_t.shape[1:], mode='bilinear')(modulator)
-    modulator = modulator.squeeze(1).detach().permute(1, 2, 0).cpu().numpy()
+    modulator = modulator.squeeze(1).detach().permute(1, 2, 0).numpy()
     modulator = (modulator - modulator.min()) / (modulator.max() - modulator.min())
     cam0 = show_cam_on_image(img_d, modulator, use_rgb=True)
 
     modulator = model.layers[0].blocks[5].modulation.modulator.norm(2, 1, keepdim=True)
     modulator = nn.Upsample(size=img_t.shape[1:], mode='bilinear')(modulator)
-    modulator = modulator.squeeze(1).detach().permute(1, 2, 0).cpu().numpy()
+    modulator = modulator.squeeze(1).detach().permute(1, 2, 0).numpy()
     modulator = (modulator - modulator.min()) / (modulator.max() - modulator.min())
     cam1 = show_cam_on_image(img_d, modulator, use_rgb=True)
 
     modulator = model.layers[0].blocks[8].modulation.modulator.norm(2, 1, keepdim=True)
     modulator = nn.Upsample(size=img_t.shape[1:], mode='bilinear')(modulator)
-    modulator = modulator.squeeze(1).detach().permute(1, 2, 0).cpu().numpy()
+    modulator = modulator.squeeze(1).detach().permute(1, 2, 0).numpy()
     modulator = (modulator - modulator.min()) / (modulator.max() - modulator.min())
     cam2 = show_cam_on_image(img_d, modulator, use_rgb=True)
 
     modulator = model.layers[0].blocks[11].modulation.modulator.norm(2, 1, keepdim=True)
     modulator = nn.Upsample(size=img_t.shape[1:], mode='bilinear')(modulator)
-    modulator = modulator.squeeze(1).detach().permute(1, 2, 0).cpu().numpy()
+    modulator = modulator.squeeze(1).detach().permute(1, 2, 0).numpy()
     modulator = (modulator - modulator.min()) / (modulator.max() - modulator.min())
     cam3 = show_cam_on_image(img_d, modulator, use_rgb=True)