Update app.py (#2)

- Update app.py (d3d0d9f7af6ab465ec836d4861e8a615fd26a9fe)


Co-authored-by: Syeda Reja <syeda-Rija20@users.noreply.huggingface.co>

app.py

@@ -464,10 +464,20 @@ def load_text_model():
     # return model
 @st.cache_resource(show_spinner=False)
 def load_image_model():
-    import os
-    import tempfile
+    # ── Patch incompatible Keras layers BEFORE any loading ──
+    try:
+        import keras
+        from keras.layers import RandomFlip, RandomRotation, RandomZoom, RandomBrightness, RandomContrast
+
+        for cls in [RandomFlip, RandomRotation, RandomZoom, RandomBrightness, RandomContrast]:
+            _orig = cls.__init__
+            def _patched(self, *args, _orig=_orig, **kwargs):
+                kwargs.pop("data_format", None)
+                _orig(self, *args, **kwargs)
+            cls.__init__ = _patched
+    except Exception:
+        pass  # If patching fails, try loading anyway
 
-    # Download the weights file
     filenames_to_try = [
         "image_detector_v2.keras",
         "image_detector_v2.h5",
@@ -486,58 +496,29 @@ def load_image_model():
 
     if path is None:
         raise RuntimeError(
-            "Could not download image model from HuggingFace Hub."
+            "Could not download image model from HuggingFace Hub. "
+            "Check that 'syeda-Rija20/image-detector' is public and the file exists."
         )
 
-    # ── Rebuild EfficientNetB3 architecture manually ──
-    # This bypasses the Keras version mismatch on RandomFlip data_format
     try:
         import keras
-        from keras import layers
-
-        # Build the same architecture the model was trained with,
-        # WITHOUT the augmentation Sequential (not needed at inference)
-        base = keras.applications.EfficientNetB3(
-            include_top=False,
-            weights=None,           # no pretrained weights; we load them below
-            input_shape=(224, 224, 3),
-        )
-        base.trainable = True
-
-        inputs = keras.Input(shape=(224, 224, 3))
-        x = base(inputs, training=False)
-        x = layers.GlobalAveragePooling2D()(x)
-        x = layers.BatchNormalization(momentum=0.99, epsilon=0.001)(x)
-        x = layers.Dense(256, activation="relu")(x)
-        x = layers.Dropout(0.4)(x, training=False)
-        outputs = layers.Dense(1, activation="sigmoid")(x)
-
-        model = keras.Model(inputs, outputs)
-        model.load_weights(path)
+        model = keras.saving.load_model(path, compile=False)
         return model
-
-    except Exception as e_rebuild:
-        # ── Last resort: patch RandomFlip then load ──
+    except Exception as e1:
         try:
-            import keras
-            from keras.layers import RandomFlip as _RF
-
-            _orig_init = _RF.__init__
-
-            def _patched_init(self, mode="horizontal", seed=None, **kwargs):
-                kwargs.pop("data_format", None)   # strip unsupported kwarg
-                _orig_init(self, mode=mode, seed=seed, **kwargs)
-
-            _RF.__init__ = _patched_init
-
-            model = keras.saving.load_model(path, compile=False)
+            import tensorflow as tf
+            model = tf.keras.models.load_model(path, compile=False)
             return model
-        except Exception as e_patch:
+        except Exception as e2:
             raise RuntimeError(
                 f"Failed to load image model.\n"
-                f"Rebuild attempt: {e_rebuild}\n"
-                f"Patch attempt: {e_patch}"
-            )    
+                f"keras attempt: {e1}\n"
+                f"tf.keras attempt: {e2}"
+            )
+
+       
+
+          
 
 
 # ─────────────────────────────────────────────