Fix bug

app.py

@@ -30,7 +30,7 @@ image = None
 ###################################
 # Functions.
 ###################################
-# @st.cache_resource
+# @st.cache_resource  # TODO: this doesn't work on Huggingface!
 def load_model(model_path: str, device: torch.device) -> BasePredictor:
     model = utils.load_is_model(model_path, device, cpu_dist_maps=True)
     predictor_params = {"brs_mode": "NoBRS"}
@@ -54,9 +54,7 @@ def feed_clicks(
         clicker.add_click(click)
 
 
-def predict(
-    image: Image, mask: torch.Tensor, threshold: float = 0.5
-) -> torch.Tensor:
+def predict(image: Image, mask: torch.Tensor, threshold: float = 0.5) -> torch.Tensor:
     predictor.set_input_image(np.array(image))
     with st.spinner("Wait for prediction..."):
         pred = predictor.get_prediction(clicker, prev_mask=mask)
@@ -120,7 +118,7 @@ if canvas_result.json_data and canvas_result.json_data["objects"] and image:
     feed_clicks(clicker, canvas_result.json_data["objects"], image_width, image_height)
 
     # Run prediction.
-    mask = torch.zeros((1, 1, image_width, image_height), device=device)
+    mask = torch.zeros((1, 1, image_height, image_width), device=device)
     pred = predict(image, mask, threshold)
 
     # Show the prediction result.