Update app.py

app.py

@@ -16,8 +16,7 @@ from io import BytesIO
 import re
 
 from surya.models.helio_spectformer import HelioSpectFormer
-# *** FIX: Import ScalerCollection and correct the import for inverse_transform_single_channel ***
-from surya.utils.data import build_scalers, ScalerCollection
+from surya.utils.data import build_scalers
 from surya.datasets.helio import inverse_transform_single_channel
 
 warnings.filterwarnings("ignore", category=UserWarning, module='sunpy')
@@ -49,10 +48,7 @@ def setup_and_load_model():
         config = yaml.safe_load(fp)
     APP_CACHE["config"] = config
     scalers_info = yaml.safe_load(open("data/Surya-1.0/scalers.yaml", "r"))
-    
-    # *** FIX: Create the ScalerCollection object from the dictionary returned by build_scalers ***
-    scalers_dict = build_scalers(info=scalers_info)
-    APP_CACHE["scalers"] = ScalerCollection(scalers_dict, config["data"]["sdo_channels"])
+    APP_CACHE["scalers"] = build_scalers(info=scalers_info)
     
     yield "Initializing model architecture..."
     model_config = config["model"]
@@ -142,7 +138,7 @@ def fetch_and_process_sdo_data(target_dt, forecast_horizon_minutes):
             images[t][channel] = Image.open(BytesIO(response.content))
             
     yield "✅ All images found and downloaded. Starting preprocessing..."
-    scaler = APP_CACHE["scalers"]
+    scalers_dict = APP_CACHE["scalers"]
     processed_tensors = {}
     for t, channel_images in images.items():
         channel_tensors = []
@@ -154,7 +150,9 @@ def fetch_and_process_sdo_data(target_dt, forecast_horizon_minutes):
             img_resized = img.resize((img_size, img_size), Image.Resampling.LANCZOS)
             norm_data = np.array(img_resized, dtype=np.float32)
 
-            scaled_data = scaler.transform(norm_data, c_idx=i)
+            # *** FIX: Retrieve the correct scaler object from the dictionary for the current channel ***
+            scaler = scalers_dict[channel]
+            scaled_data = scaler.transform(norm_data.reshape(-1, 1)).reshape(norm_data.shape)
             channel_tensors.append(torch.from_numpy(scaled_data.astype(np.float32)))
         processed_tensors[t] = torch.stack(channel_tensors)
 
@@ -180,13 +178,17 @@ def run_inference(input_tensor):
 def generate_visualization(last_input_map, prediction_tensor, target_map, channel_name):
     if last_input_map is None: return None, None, None
     c_idx = SDO_CHANNELS.index(channel_name)
-    scaler = APP_CACHE["scalers"]
-    all_means, all_stds, all_epsilons, all_sl_scale_factors = scaler.get_params()
-    mean, std, epsilon, sl_scale_factor = all_means[c_idx], all_stds[c_idx], all_epsilons[c_idx], all_sl_scale_factors[c_idx]
-    pred_slice = inverse_transform_single_channel(
-        prediction_tensor[0, c_idx].numpy(), mean=mean, std=std, epsilon=epsilon, sl_scale_factor=sl_scale_factor
-    )
     
+    # *** FIX: Retrieve the correct scaler object for the current channel to get its parameters ***
+    scaler = APP_CACHE["scalers"][channel_name]
+    params = scaler.to_dict()
+    mean, std = params['mean'], params['std']
+    
+    # Note: The inverse transform for the simplified JPEG pipeline might differ from the original
+    # We will use a standard inverse scaling, which is the most logical approach here.
+    pred_slice_scaled = prediction_tensor[0, c_idx].numpy()
+    pred_slice = (pred_slice_scaled * std) + mean
+
     target_img_data = np.array(target_map[channel_name])
     vmax = np.quantile(np.nan_to_num(target_img_data), 0.995)
     cmap_name = f"sdoaia{channel_name.replace('aia', '')}" if 'aia' in channel_name else 'hmimag'