Update app.py

app.py

@@ -4,12 +4,7 @@ from huggingface_hub import snapshot_download
 import yaml
 import numpy as np
 from PIL import Image
-import sunpy.map
-import sunpy.net.attrs as a
-from sunpy.net import Fido
-from astropy.wcs import WCS
-import astropy.units as u
-from reproject import reproject_interp
+import requests
 import os
 import warnings
 import logging
@@ -17,6 +12,7 @@ import datetime
 import matplotlib.pyplot as plt
 import sunpy.visualization.colormaps as sunpy_cm
 import traceback
+from io import BytesIO
 
 from surya.models.helio_spectformer import HelioSpectFormer
 from surya.utils.data import build_scalers
@@ -29,22 +25,13 @@ logger = logging.getLogger(__name__)
 
 APP_CACHE = {}
 
-SDO_CHANNELS_MAP = {
-    "aia94": (a.Wavelength(94 * u.angstrom), a.Sample(12 * u.s)),
-    "aia131": (a.Wavelength(131 * u.angstrom), a.Sample(12 * u.s)),
-    "aia171": (a.Wavelength(171 * u.angstrom), a.Sample(12 * u.s)),
-    "aia193": (a.Wavelength(193 * u.angstrom), a.Sample(12 * u.s)),
-    "aia211": (a.Wavelength(211 * u.angstrom), a.Sample(12 * u.s)),
-    "aia304": (a.Wavelength(304 * u.angstrom), a.Sample(12 * u.s)),
-    "aia335": (a.Wavelength(335 * u.angstrom), a.Sample(12 * u.s)),
-    "aia1600": (a.Wavelength(1600 * u.angstrom), a.Sample(24 * u.s)),
-    "hmi_m": (a.Physobs("intensity"), a.Sample(45 * u.s)),
-    "hmi_bx": (a.Physobs("los_magnetic_field"), a.Sample(720 * u.s)),
-    "hmi_by": (a.Physobs("los_magnetic_field"), a.Sample(720 * u.s)),
-    "hmi_bz": (a.Physobs("los_magnetic_field"), a.Sample(720 * u.s)),
-    "hmi_v": (a.Physobs("los_velocity"), a.Sample(45 * u.s)),
+CHANNEL_TO_URL_CODE = {
+    "aia94": "0094", "aia131": "0131", "aia171": "0171", "aia193": "0193",
+    "aia211": "0211", "aia304": "0304", "aia335": "0335", "aia1600": "1600",
+    "hmi_m": "HMIBC", "hmi_bx": "HMIB", "hmi_by": "HMIB", 
+    "hmi_bz": "HMIB", "hmi_v": "HMID"
 }
-SDO_CHANNELS = list(SDO_CHANNELS_MAP.keys())
+SDO_CHANNELS = list(CHANNEL_TO_URL_CODE.keys())
 
 def setup_and_load_model():
     if "model" in APP_CACHE:
@@ -87,67 +74,56 @@ def setup_and_load_model():
     APP_CACHE["model"] = model
     yield "✅ Model setup complete."
 
+def fetch_browse_image(channel, target_dt, max_retries=15):
+    url_code = CHANNEL_TO_URL_CODE[channel]
+    base_url = "https://sdo.gsfc.nasa.gov/assets/img/browse"
+    
+    for i in range(max_retries):
+        dt_to_try = target_dt - datetime.timedelta(minutes=i)
+        date_str = dt_to_try.strftime("%Y/%m/%d")
+        img_str = dt_to_try.strftime(f"%Y%m%d_%H%M%S_4096_{url_code}.jpg")
+        url = f"{base_url}/{date_str}/{img_str}"
+        
+        response = requests.get(url)
+        if response.status_code == 200:
+            logger.info(f"Successfully found image for {channel} at {dt_to_try}")
+            return Image.open(BytesIO(response.content))
+    
+    raise FileNotFoundError(f"Could not find any recent image for {channel} within {max_retries} minutes of {target_dt}.")
+
 def fetch_and_process_sdo_data(target_dt, forecast_horizon_minutes):
     config = APP_CACHE["config"]
     img_size = config["model"]["img_size"]
     
     input_deltas = config["data"]["time_delta_input_minutes"]
-    target_delta = forecast_horizon_minutes
     input_times = [target_dt + datetime.timedelta(minutes=m) for m in input_deltas]
-    target_time = target_dt + datetime.timedelta(minutes=target_delta)
+    target_time = target_dt + datetime.timedelta(minutes=forecast_horizon_minutes)
     all_times = sorted(list(set(input_times + [target_time])))
 
-    data_maps = {}
-    last_successful_map = {}
-    total_downloads = len(all_times) * len(SDO_CHANNELS)
-    downloads_done = 0
-    yield f"Starting download of {total_downloads} data files..."
+    images = {}
+    total_fetches = len(all_times) * len(SDO_CHANNELS)
+    fetches_done = 0
+    yield f"Starting search for {total_fetches} data files..."
+
     for t in all_times:
-        data_maps[t] = {}
-        for i, (channel, (physobs, sample)) in enumerate(SDO_CHANNELS_MAP.items()):
-            downloads_done += 1
-            yield f"Downloading [{downloads_done}/{total_downloads}]: {channel} for {t.strftime('%Y-%m-%d %H:%M')}..."
-            
-            if channel in ["hmi_by", "hmi_bz"]: 
-                if data_maps[t].get("hmi_bx"):
-                    smap = data_maps[t]["hmi_bx"]
-                    data_maps[t][channel] = smap
-                    last_successful_map[channel] = smap
-                continue
+        images[t] = {}
+        for channel in SDO_CHANNELS:
+            fetches_done += 1
+            yield f"Searching [{fetches_done}/{total_fetches}]: {channel} near {t.strftime('%Y-%m-%d %H:%M')}..."
+            images[t][channel] = fetch_browse_image(channel, t)
             
-            time_attr = a.Time(t - datetime.timedelta(minutes=5), t + datetime.timedelta(minutes=5))
-            instrument = a.Instrument.hmi if "hmi" in channel else a.Instrument.aia
-            query = Fido.search(time_attr, instrument, physobs, sample)
-            
-            if query:
-                files = Fido.fetch(query[0,0], path="./data/sdo_cache")
-                smap = sunpy.map.Map(files[0])
-                data_maps[t][channel] = smap
-                last_successful_map[channel] = smap
-            elif channel in last_successful_map:
-                yield f"⚠️ WARNING: No data for {channel} near {t}. Reusing previous image."
-                data_maps[t][channel] = last_successful_map[channel]
-            else:
-                raise ValueError(f"CRITICAL: No data found for initial image of {channel}. Cannot proceed.")
-            
-    yield "✅ All files downloaded. Starting preprocessing..."
-    output_wcs = WCS(naxis=2)
-    output_wcs.wcs.crpix = [(img_size + 1) / 2, (img_size + 1) / 2]
-    output_wcs.wcs.cdelt = np.array([-1.2, 1.2]) * u.arcsec
-    output_wcs.wcs.crval = [0, 0] * u.arcsec
-    output_wcs.wcs.ctype = ['HPLN-TAN', 'HPLT-TAN']
-
+    yield "✅ All images found. Starting preprocessing..."
     scaler = APP_CACHE["scalers"]
     processed_tensors = {}
-    for t, channel_maps in data_maps.items():
+    for t, channel_images in images.items():
         channel_tensors = []
         for i, channel in enumerate(SDO_CHANNELS):
-            smap = channel_maps[channel]
-            reprojected_data, _ = reproject_interp(smap, output_wcs, shape_out=(img_size, img_size))
+            img = channel_images[channel]
+            if img.mode != 'L':
+                img = img.convert('L')
             
-            exp_time = smap.meta.get('exptime', 1.0)
-            if exp_time is None or exp_time <= 0: exp_time = 1.0
-            norm_data = reprojected_data / exp_time
+            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.reshape(-1, 1), c_idx=i).reshape(norm_data.shape)
             channel_tensors.append(torch.from_numpy(scaled_data.astype(np.float32)))
@@ -156,10 +132,10 @@ def fetch_and_process_sdo_data(target_dt, forecast_horizon_minutes):
     yield "✅ Preprocessing complete."
     input_tensor_list = [processed_tensors[t] for t in input_times]
     input_tensor = torch.stack(input_tensor_list, dim=1).unsqueeze(0)
-    target_map = data_maps[target_time]
-    last_input_map = data_maps[input_times[-1]]
+    target_image_map = images[target_time]
+    last_input_image_map = images[input_times[-1]]
 
-    yield (input_tensor, last_input_map, target_map)
+    yield (input_tensor, last_input_image_map, target_image_map)
 
 def run_inference(input_tensor):
     model = APP_CACHE["model"]
@@ -181,17 +157,20 @@ def generate_visualization(last_input_map, prediction_tensor, target_map, channe
     pred_slice = inverse_transform_single_channel(
         prediction_tensor[0, c_idx].numpy(), mean=mean, std=std, epsilon=epsilon, sl_scale_factor=sl_scale_factor
     )
-    vmax = np.quantile(np.nan_to_num(target_map[channel_name].data), 0.995)
+    
+    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'
     cmap = plt.get_cmap(sunpy_cm.cmlist.get(cmap_name, 'gray'))
-    def to_pil(data, flip=False):
+    
+    def to_pil(data):
         data_clipped = np.nan_to_num(data)
         data_clipped = np.clip(data_clipped, 0, vmax)
         data_norm = data_clipped / vmax if vmax > 0 else data_clipped
         colored = (cmap(data_norm)[:, :, :3] * 255).astype(np.uint8)
-        img = Image.fromarray(colored)
-        return img.transpose(Image.Transpose.FLIP_TOP_BOTTOM) if flip else img
-    return to_pil(last_input_map[channel_name].data, flip=True), to_pil(pred_slice, flip=True), to_pil(target_map[channel_name].data, flip=True)
+        return Image.fromarray(colored)
+        
+    return last_input_map[channel_name], to_pil(pred_slice), target_map[channel_name]
 
 def forecast_controller(date_str, hour, minute, forecast_horizon):
     yield {
@@ -265,7 +244,9 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         # ☀️ Surya: Live Forecast Demo ☀️
         ### A Foundation Model for Solar Dynamics
         This demo runs NASA's **Surya**, a foundation model trained to understand the physics of the Sun.
-        It looks at the Sun in 13 different channels (8 from the AIA instrument, 5 from HMI) simultaneously to learn the complex relationships between solar phenomena like coronal loops, magnetic fields, and solar flares. By seeing these interconnected views, it can generate a holistic forecast of what the entire solar disk will look like in the near future.
+        It looks at the Sun in 13 different channels (wavelengths of light) simultaneously to learn the complex relationships between phenomena like coronal loops, magnetic fields, and solar flares. By seeing these interconnected views, it can generate a holistic forecast of what the entire solar disk will look like in the near future.
+        <br>
+        <p style="color:red;font-weight:bold;">NOTE: This demo uses lower-quality browse images for reliability. The model was trained on high-fidelity scientific data, so forecast accuracy may vary.</p>
         </div>
         """
     )
@@ -273,11 +254,11 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     with gr.Accordion("Step 1: Configure Forecast", open=True):
         with gr.Row():
             date_input = gr.Textbox(
-                label="Date",
-                value=datetime.date.today().strftime("%Y-%m-%d")
+                label="Date (YYYY-MM-DD)",
+                value=(datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(hours=3)).strftime("%Y-%m-%d")
             )
-            hour_slider = gr.Slider(label="Hour (UTC)", minimum=0, maximum=23, step=1, value=datetime.datetime.utcnow().hour - 3)
-            minute_slider = gr.Slider(label="Minute", minimum=0, maximum=59, step=1, value=datetime.datetime.utcnow().minute)
+            hour_slider = gr.Slider(label="Hour (UTC)", minimum=0, maximum=23, step=1, value=(datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(hours=3)).hour)
+            minute_slider = gr.Slider(label="Minute (UTC)", minimum=0, maximum=59, step=1, value=(datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(hours=3)).minute)
         horizon_slider = gr.Slider(
             label="Forecast Horizon (minutes ahead)",
             minimum=12, maximum=120, step=12, value=12
@@ -285,7 +266,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     
     run_button = gr.Button("🔮 Generate Forecast", variant="primary")
     
-    with gr.Accordion("Step 2: View Log", open=False) as log_accordion:
+    with gr.Accordion("Step 2: View Log", open=False):
         log_box = gr.Textbox(label="Log", interactive=False, visible=False, lines=5, max_lines=10)
 
     with gr.Group(visible=False) as results_group:
@@ -315,5 +296,4 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     )
 
 if __name__ == "__main__":
-    os.makedirs("./data/sdo_cache", exist_ok=True)
     demo.launch(debug=True)