Fix app

app.py

@@ -22,48 +22,6 @@ os.system(f'cp {model_inference} .')
 
 from inference import process_channel_group, _convert_np_uint8, load_example, run_model
 
-def extract_rgb_imgs(input_img, pred_img, channels):
-    """ Wrapper function to save Geotiff images (original, reconstructed, masked) per timestamp.
-    Args:
-        input_img: input torch.Tensor with shape (C, H, W).
-        rec_img: reconstructed torch.Tensor with shape (C, T, H, W).
-        pred_img: mask torch.Tensor with shape (C, T, H, W).
-        channels: list of indices representing RGB channels.
-        mean: list of mean values for each band.
-        std: list of std values for each band.
-        output_dir: directory where to save outputs.
-        meta_data: list of dicts with geotiff meta info.
-    """
-    rgb_orig_list = []
-    rgb_mask_list = []
-    rgb_pred_list = []
-
-    for t in range(input_img.shape[1]):
-        rgb_orig, rgb_pred = process_channel_group(orig_img=input_img[:, t, :, :],
-                                                   new_img=rec_img[:, t, :, :],
-                                                   channels=channels,
-                                                   mean=mean,
-                                                   std=std)
-
-        rgb_mask = mask_img[channels, t, :, :] * rgb_orig
-
-        # extract images
-        rgb_orig_list.append(_convert_np_uint8(rgb_orig).transpose(1, 2, 0))
-        rgb_mask_list.append(_convert_np_uint8(rgb_mask).transpose(1, 2, 0))
-        rgb_pred_list.append(_convert_np_uint8(rgb_pred).transpose(1, 2, 0))
-
-    # Add white dummy image values for missing timestamps
-    dummy = np.ones((20, 20), dtype=np.uint8) * 255
-    num_dummies = 4 - len(rgb_orig_list)
-    if num_dummies:
-        rgb_orig_list.extend([dummy] * num_dummies)
-        rgb_mask_list.extend([dummy] * num_dummies)
-        rgb_pred_list.extend([dummy] * num_dummies)
-
-    outputs = rgb_orig_list + rgb_mask_list + rgb_pred_list
-
-    return outputs
-
 
 def predict_on_images(data_file: str | Path, config_path: str, checkpoint: str):
     try:
@@ -81,7 +39,7 @@ def predict_on_images(data_file: str | Path, config_path: str, checkpoint: str):
     # Load model ---------------------------------------------------------------------------------
 
     lightning_model = LightningInferenceModel.from_config(config_path, checkpoint)
-    img_size = 256  # Size of Sen1Floods11
+    img_size = 512  # Size from Sen1Floods11 training
 
     # Loading data ---------------------------------------------------------------------------------
 

requirements.txt

@@ -5,4 +5,4 @@ rasterio
 einops
 huggingface_hub
 gradio
-git+https://github.com/IBM/terratorch.git
+terratorch==1.0.2