Upload app.py

app.py

@@ -782,7 +782,7 @@ with tab2:
 with tab3:
     st.header("📂 Load Pre-trained Models")
 
-    st.markdown("Upload Bayesian Network models (.pkl files) to view their structures.")
+    st.markdown("Upload Bayesian Network models (.pkl or .json files) to view their structures.")
     st.markdown("**Maximum: 2 models**")
     st.markdown("---")
 
@@ -792,16 +792,17 @@ with tab3:
         uploaded_model = None
     else:
         uploaded_model = st.file_uploader(
-            "Upload model file (.pkl)",
-            type=['pkl'],
-            help="Upload a .pkl file containing a Bayesian Network model"
+            "Upload model file (.pkl or .json)",
+            type=['pkl', 'json'],
+            help="Upload a .pkl or .json file containing a Bayesian Network model"
         )
 
     if uploaded_model:
         col_load1, col_load2 = st.columns([3, 1])
 
         with col_load1:
-            st.info(f"📄 File: **{uploaded_model.name}**")
+            file_type = "PKL" if uploaded_model.name.endswith('.pkl') else "JSON"
+            st.info(f"📄 File: **{uploaded_model.name}** ({file_type})")
 
         with col_load2:
             if st.button("🔄 Load Model", type="primary", use_container_width=True):
@@ -810,28 +811,68 @@ with tab3:
                         import pickle
                         import tempfile
                         import os
+                        import json
+                        from pgmpy.models import BayesianNetwork
+                        from pgmpy.factors.discrete import TabularCPD
 
-                        # Save to temp file
-                        with tempfile.NamedTemporaryFile(delete=False, suffix='.pkl') as tmp_file:
-                            tmp_file.write(uploaded_model.read())
-                            tmp_path = tmp_file.name
+                        model = None
 
-                        # Load model
-                        with open(tmp_path, 'rb') as f:
-                            model_data = pickle.load(f)
+                        # Check file type
+                        if uploaded_model.name.endswith('.pkl'):
+                            # Load PKL file
+                            with tempfile.NamedTemporaryFile(delete=False, suffix='.pkl') as tmp_file:
+                                tmp_file.write(uploaded_model.read())
+                                tmp_path = tmp_file.name
 
-                        os.unlink(tmp_path)
+                            with open(tmp_path, 'rb') as f:
+                                model_data = pickle.load(f)
 
-                        # Extract model
-                        from pgmpy.models import BayesianNetwork
+                            os.unlink(tmp_path)
 
-                        if isinstance(model_data, BayesianNetwork):
-                            model = model_data
-                        elif isinstance(model_data, dict) and 'model' in model_data:
-                            model = model_data['model']
-                        else:
-                            st.error("❌ Invalid model format")
-                            model = None
+                            if isinstance(model_data, BayesianNetwork):
+                                model = model_data
+                            elif isinstance(model_data, dict) and 'model' in model_data:
+                                model = model_data['model']
+                            else:
+                                st.error("❌ Invalid PKL model format")
+
+                        elif uploaded_model.name.endswith('.json'):
+                            # Load JSON file
+                            json_data = json.loads(uploaded_model.read().decode('utf-8'))
+
+                            # Reconstruct BayesianNetwork from JSON
+                            if 'model' in json_data:
+                                model_data = json_data['model']
+
+                                # Create network with edges
+                                model = BayesianNetwork(model_data['edges'])
+
+                                # Add CPDs
+                                for node, cpd_data in model_data['cpds'].items():
+                                    # Ensure values is 2D array
+                                    import numpy as np
+                                    values = cpd_data['values']
+                                    if isinstance(values, list) and len(values) > 0:
+                                        if not isinstance(values[0], list):
+                                            # 1D array -> convert to 2D
+                                            values = np.array(values).reshape(-1, 1)
+
+                                    cpd = TabularCPD(
+                                        variable=cpd_data['variable'],
+                                        variable_card=cpd_data['cardinality'][0],
+                                        values=values,
+                                        evidence=cpd_data['parents'] or None,
+                                        evidence_card=cpd_data['cardinality'][1:] or None,
+                                        state_names=cpd_data.get('state_names', None)
+                                    )
+                                    model.add_cpds(cpd)
+
+                                # Validate model
+                                if not model.check_model():
+                                    st.error("❌ Invalid model structure")
+                                    model = None
+                            else:
+                                st.error("❌ Invalid JSON format: missing 'model' key")
 
                         if model:
                             # Store model