Streamlit: fix image arg and indentation; integrate Imagen override rule into final decision

streamlit_app.py

@@ -76,8 +76,8 @@ def main():
             st.error("Please provide a valid image (upload or URL).")
             return
 
-        # Show the image
-        st.image(tmp_path, caption="Input Image", use_container_width=True)
+        # Show the image (use_column_width for Streamlit 1.x)
+        st.image(tmp_path, caption="Input Image", use_column_width=True)
 
         st.info("Loading models (first time may take a bit)...")
         cascade = get_cascade()

two_stage_inference.py

@@ -293,16 +293,38 @@ class CascadeClassifier:
         m_conf, m_idx, m_probs = softmax_confidence(logits)
         m_label = self.multi_id2label.get(int(m_idx), str(m_idx))
 
+        # Apply Imagen override rule for the final decision:
+        # If top-1 is 'Imagen' and any other class prob > 1.7e-5, choose the second-highest instead.
+        override_idx = int(m_idx)
+        override_conf = float(m_conf)
+        try:
+            # Find index for label 'Imagen' (case-insensitive)
+            imagen_idx = next((i for i, v in self.multi_id2label.items() if str(v).strip().lower() == 'imagen'), None)
+        except Exception:
+            imagen_idx = None
+        if imagen_idx is not None:
+            # Sort indices by probability descending
+            order = sorted(range(len(m_probs)), key=lambda i: m_probs[i], reverse=True)
+            if order and order[0] == int(imagen_idx):
+                # Check if any non-imagen prob exceeds threshold
+                THRESH = 1.7e-5
+                any_other_above = any((i != imagen_idx) and (m_probs[i] > THRESH) for i in range(len(m_probs)))
+                if any_other_above and len(order) > 1:
+                    override_idx = int(order[1])
+                    override_conf = float(m_probs[override_idx])
+
         result['multiclass'] = {
             'label': m_label,
             'index': int(m_idx),
             'confidence': float(m_conf),
             'probs': m_probs,
         }
+        # Final decision reflects the override rule if triggered
+        final_label = self.multi_id2label.get(int(override_idx), str(override_idx))
         result['final'] = {
             'decision': 'fake',
-            'origin': m_label,
-            'confidence': float(m_conf),
+            'origin': final_label,
+            'confidence': float(override_conf),
         }
         return result