Update app.py

Fix normailzation for SAT images

app.py

@@ -82,6 +82,18 @@ DEFAULT_DATASET_LABEL = "LVD-1689M"  # initial radio
 DEFAULT_OVERLAY_ALPHA = 0.55
 DEFAULT_SHOW_GRID = True
 
+# ---------- Normalization presets ----------
+NORMALIZE_STATS = {
+    "lvd1689m": {
+        "mean": [0.485, 0.456, 0.406],
+        "std":  [0.229, 0.224, 0.225],
+    },
+    "sat493m": {
+        "mean": [0.430, 0.411, 0.296],
+        "std":  [0.213, 0.156, 0.143],
+    },
+}
+
 # ---------- Sample image URLs (dependent on dataset) ----------
 SAMPLE_URL_CHOICES: Dict[str, List[Tuple[str, str]]] = {
     # LVD: current ones
@@ -135,17 +147,22 @@ def pad_to_multiple(pil_img: Image.Image, multiple: int = 16) -> Tuple[Image.Ima
     canvas.paste(pil_img, (0, 0))
     return canvas, (0, 0, W_pad - W, H_pad - H)
 
-def preprocess_no_resize(pil_img: Image.Image, multiple: int = 16):
+def preprocess_no_resize(pil_img: Image.Image, multiple: int = 16, dataset_key: str = "lvd1689m"):
     img_padded, pad_box = pad_to_multiple(pil_img, multiple=multiple)
+
+    # Pick stats based on dataset (default to LVD if unknown)
+    stats = NORMALIZE_STATS.get(dataset_key, NORMALIZE_STATS["lvd1689m"])
+    mean, std = stats["mean"], stats["std"]
+
     transform = transforms.Compose([
         transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406],
-                             std =[0.229, 0.224, 0.225]),
+        transforms.Normalize(mean=mean, std=std),
     ])
     pixel_tensor = transform(img_padded).unsqueeze(0)  # (1,3,H,W)
     disp_np = np.array(img_padded, dtype=np.uint8)
     return {"pixel_values": pixel_tensor}, disp_np, pad_box
 
+
 def upsample_nearest(arr: np.ndarray, H: int, W: int, ps: int) -> np.ndarray:
     if arr.ndim == 2:
         return arr.repeat(ps, 0).repeat(ps, 1)
@@ -198,10 +215,11 @@ def infer_patch_size(model, default: int = 16) -> int:
 
 # ---------- Per-image state ----------
 class PatchImageState:
-    def __init__(self, pil_img: Image.Image, model, device_str: str, ps: int):
+    def __init__(self, pil_img: Image.Image, model, device_str: str, ps: int, dataset_key: str):
         self.pil = pil_img
         self.ps = ps
-        inputs, disp_np, _ = preprocess_no_resize(pil_img, multiple=ps)
+        self.dataset_key = dataset_key
+        inputs, disp_np, _ = preprocess_no_resize(pil_img, multiple=ps, dataset_key=dataset_key)
         self.disp = disp_np
         pv = inputs["pixel_values"].to(device_str)  # (1,3,H,W)
         _, _, H, W = pv.shape
@@ -344,9 +362,12 @@ def init_states(
     device_str = "cuda" if torch.cuda.is_available() else "cpu"
     model = load_model_cached(full_model_id, device_str)
     ps = infer_patch_size(model, 16)
+    
+    # Get dataset_key ("lvd1689m" or "sat493m") from the radio label
+    dataset_key = dataset_label_to_key(dataset_label)
 
-    left_state = PatchImageState(left_img, model, device_str, ps) if left_img is not None else None
-    right_state = PatchImageState(right_img, model, device_str, ps) if right_img is not None else None
+    left_state = PatchImageState(left_img, model, device_str, ps, dataset_key) if left_img is not None else None
+    right_state = PatchImageState(right_img, model, device_str, ps, dataset_key) if right_img is not None else None
 
     active_side = 0 if left_state is not None else 1