add padding

app.py

@@ -47,7 +47,7 @@ def process_channel_group(orig_img, new_img, channels, data_mean, data_std):
     for c in channels:
         orig_ch = orig_img[c, ...]
         valid_mask = torch.ones_like(orig_ch, dtype=torch.bool)
-        valid_mask[orig_ch == 0.0001] = False
+        valid_mask[orig_ch == NO_DATA_FLOAT] = False
 
         # Back to original data range
         orig_ch = (orig_ch * data_std[c]) + data_mean[c]
@@ -138,8 +138,8 @@ def load_example(file_paths: List[str], mean: List[float], std: List[float]):
         imgs.append(img)
         metas.append(meta)
 
-    imgs = np.stack(imgs, axis=0)    # num_frames, img_size, img_size, C
-    imgs = np.moveaxis(imgs, -1, 0).astype('float32')  # C, num_frames, img_size, img_size
+    imgs = np.stack(imgs, axis=0)    # num_frames, H, W, C
+    imgs = np.moveaxis(imgs, -1, 0).astype('float32')  # C, num_frames, H, W
     imgs = np.expand_dims(imgs, axis=0)  # add batch dim
 
     return imgs, metas
@@ -308,7 +308,7 @@ def predict_on_images(data_files: list, mask_ratio: float, yaml_file_path: str,
             norm_pix_loss=False)
 
     total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
-    print(f"\n--> model has {total_params / 1e6} Million params.\n")
+    print(f"\n--> Model has {total_params:,} parameters.\n")
 
     model.to(device)
 
@@ -320,6 +320,12 @@ def predict_on_images(data_files: list, mask_ratio: float, yaml_file_path: str,
 
     model.eval()
     channels = [bands.index(b) for b in ['B04', 'B03', 'B02']]  # BGR -> RGB
+    
+    # Reflect pad if not divisible by img_size
+    original_h, original_w = input_data.shape[-2:]
+    pad_h = img_size - (original_h % img_size)
+    pad_w = img_size - (original_w % img_size)
+    input_data = np.pad(input_data, ((0, 0), (0, 0), (0, 0), (0, pad_h), (0, pad_w)), mode='reflect')
 
     # Build sliding window
     batch = torch.tensor(input_data, device='cpu')
@@ -348,13 +354,10 @@ def predict_on_images(data_files: list, mask_ratio: float, yaml_file_path: str,
     mask_imgs = rearrange(mask_imgs, '(b h1 w1) c t h w -> b c t (h1 h) (w1 w)',
                           h=img_size, w=img_size, b=1, c=len(bands), t=num_frames, h1=h1, w1=w1)
 
-    # Mix original image with patches
-    h, w = rec_imgs.shape[-2:]
-    rec_imgs_full = batch.clone()
-    rec_imgs_full[..., :h, :w] = rec_imgs
-
-    mask_imgs_full = torch.ones_like(batch)
-    mask_imgs_full[..., :h, :w] = mask_imgs
+    # Cut padded images back to original size
+    rec_imgs_full = rec_imgs[..., :original_h, :original_w]
+    mask_imgs_full = mask_imgs[..., :original_h, :original_w]
+    batch_full = batch[..., :original_h, :original_w]
 
     # Build RGB images
     for d in meta_data:
@@ -363,7 +366,7 @@ def predict_on_images(data_files: list, mask_ratio: float, yaml_file_path: str,
     # save_rgb_imgs(batch[0, ...], rec_imgs_full[0, ...], mask_imgs_full[0, ...],
     #               channels, mean, std, output_dir, meta_data)
 
-    outputs = extract_rgb_imgs(batch[0, ...], rec_imgs_full[0, ...], mask_imgs_full[0, ...],
+    outputs = extract_rgb_imgs(batch_full[0, ...], rec_imgs_full[0, ...], mask_imgs_full[0, ...],
                   channels, mean, std)