Update app.py

app.py

@@ -143,26 +143,26 @@ def load_models(unet_weights_path, generator_path=None):
     generator = None
     if generator_path:
         try:
-            # Try to load the generator model
-            import tensorflow as tf
+            # Try direct loading first (works in Keras 3.x)
+            generator = tf.keras.models.load_model(generator_path)
+        except Exception as e:
+            st.warning(f"Direct model loading failed: {e}")
+            st.info("Trying alternative loading method...")
             
-            # Option 1: Try direct loading with custom objects
             try:
-                custom_objects = {'Functional': tf.keras.Model}
-                generator = tf.keras.models.load_model(generator_path, custom_objects=custom_objects)
-            except:
-                # Option 2: If that fails, try loading just the weights into a compatible architecture
-                st.warning("Attempting to recreate generator model architecture...")
+                # Create a compatible generator model
                 generator = create_pix2pix_generator()
-                generator.load_weights(generator_path, by_name=True, skip_mismatch=True)
-                st.success("Successfully loaded generator weights into a compatible model")
                 
-        except Exception as e:
-            st.error(f"Error loading generator model: {e}")
+                # Try to load weights
+                generator.load_weights(generator_path, by_name=True, skip_mismatch=True)
+                st.success("Loaded generator weights into a compatible model")
+            except Exception as e2:
+                st.error(f"Alternative loading also failed: {e2}")
     
     return unet, generator
 
 
+
 # Preprocess SAR data for SAR to Optical Translation
 def preprocess_sar_for_optical(sar_data):
     """Preprocess SAR data"""
@@ -201,16 +201,51 @@ def process_image(sar_image, unet_model, generator_model=None):
     # Get segmentation using U-Net
     seg_mask = unet_model.predict(sar_image)
     
-    # Generate optical using segmentation if generator is available
     # Generate optical using segmentation if generator is available
     colorized = None
     if generator_model:
-        colorized = generator_model.predict([sar_image, seg_mask])
-        colorized = colorized[0]
-
+        try:
+            colorized = generator_model.predict([sar_image, seg_mask])
+            colorized = colorized[0]
+            
+            # Check if output is valid
+            if np.var(colorized) < 0.01:
+                st.warning("Generator produced low-variance output. Using segmentation-based colorization instead.")
+                colorized = None
+        except Exception as e:
+            st.error(f"Error in generator prediction: {e}")
+    
+    # If generator failed or produced invalid output, create a colorized version from segmentation
+    if colorized is None and generator_model is not None:
+        st.info("Creating colorized image from segmentation mask")
+        # Create a colorized version based on segmentation
+        pred_class = np.argmax(seg_mask[0], axis=-1)
+        colorized = np.zeros((pred_class.shape[0], pred_class.shape[1], 3), dtype=np.float32)
+        
+        colors = {
+            0: [0/255, 100/255, 0/255],     # Trees - Dark green
+            1: [255/255, 165/255, 0/255],   # Shrubland - Orange
+            2: [144/255, 238/255, 144/255], # Grassland - Light green
+            3: [255/255, 255/255, 0/255],   # Cropland - Yellow
+            4: [255/255, 0/255, 0/255],     # Built-up - Red
+            5: [139/255, 69/255, 19/255],   # Bare - Brown
+            6: [255/255, 255/255, 255/255], # Snow - White
+            7: [0/255, 0/255, 255/255],     # Water - Blue
+            8: [0/255, 139/255, 139/255],   # Wetland - Dark cyan
+            9: [0/255, 255/255, 0/255],     # Mangroves - Bright green
+            10: [220/255, 220/255, 220/255] # Moss - Light grey
+        }
+        
+        for class_idx, color in colors.items():
+            colorized[pred_class == class_idx] = color
+        
+        # Convert to tanh range (-1 to 1) to match generator output format
+        colorized = colorized * 2 - 1
+    
     return seg_mask[0], colorized
 
 
+
 # Visualize results for SAR to Optical Translation
 def visualize_results(sar_image, seg_mask, colorized=None):
     # ESA WorldCover colors