added legend for kg visualization

ProtHGT_app.py

@@ -630,7 +630,9 @@ if st.session_state.submitted:
                             with open(viz_info['path'], 'r', encoding='utf-8') as f:
                                 html_content = f.read()
 
-                            st.components.v1.html(html_content, height=600)
+                            st.components.v1.html(html_content, height=1200)
+
+
 
 
             else:

visualize_kg.py

@@ -180,7 +180,7 @@ def visualize_protein_subgraph(data, protein_id, prediction_df, limit=10):
             "multiselect": true
         },
         "configure": {
-            "enabled": true,
+            "enabled": false,
             "filter": ["physics", "layout", "manipulation"],
             "showButton": true
         },
@@ -265,9 +265,103 @@ def visualize_protein_subgraph(data, protein_id, prediction_df, limit=10):
                         length=200,
                         smooth={'type': 'curvedCW', 'roundness': 0.1})
 
+    # LEGEND
+    legend_html = """
+    <style>
+        .kg-legend {
+            margin-top: 20px;
+            padding: 20px;
+            border: 1px solid #ddd;
+            border-radius: 5px;
+            font-family: Arial, sans-serif;
+        }
+        .legend-section {
+            margin-bottom: 20px;
+        }
+        .legend-title {
+            margin-bottom: 15px;
+            color: #333;
+            font-size: 16px;
+            font-weight: bold;
+        }
+        .legend-grid {
+            display: grid;
+            grid-template-columns: repeat(auto-fit, minmax(180px, 1fr));
+            gap: 12px;
+        }
+        .legend-item {
+            display: flex;
+            align-items: center;
+            padding: 4px;
+        }
+        .node-indicator {
+            width: 15px;
+            height: 15px;
+            border-radius: 50%;
+            margin-right: 10px;
+            flex-shrink: 0;
+        }
+        .edge-indicator {
+            width: 40px;
+            height: 3px;
+            margin-right: 10px;
+            flex-shrink: 0;
+        }
+        .legend-label {
+            font-size: 14px;
+        }
+    </style>
+    <div class="kg-legend">
+        <div class="legend-section">
+            <div class="legend-title">Node Types</div>
+            <div class="legend-grid">"""
+
+    # Node types
+    for node_type, color in NODE_TYPE_COLORS.items():
+        legend_html += f"""
+                <div class="legend-item">
+                    <div class="node-indicator" style="background-color: {color};"></div>
+                    <span class="legend-label">{node_type}</span>
+                </div>"""
+
+    # Edge types
+    legend_html += """
+            </div>
+        </div>
+        <div class="legend-section">
+            <div class="legend-title">Edge Colors</div>
+            <div class="legend-grid">
+                <div class="legend-item">
+                    <div class="edge-indicator" style="background-color: #8338ec;"></div>
+                    <span class="legend-label">Validated GO Prediction</span>
+                </div>
+                <div class="legend-item">
+                    <div class="edge-indicator" style="background-color: #c1121f;"></div>
+                    <span class="legend-label">Non-validated GO Prediction</span>
+                </div>
+                <div class="legend-item">
+                    <div class="edge-indicator" style="background-color: #219ebc;"></div>
+                    <span class="legend-label">Ground Truth GO Annotation</span>
+                </div>
+                <div class="legend-item">
+                    <div class="edge-indicator" style="background-color: #666666;"></div>
+                    <span class="legend-label">Other Relationships</span>
+                </div>
+            </div>
+        </div>
+    </div>
+    """
+
     # Save graph to a protein-specific file in a temporary directory
     os.makedirs('temp_viz', exist_ok=True)
     file_path = os.path.join('temp_viz', f'{protein_id}_graph.html')
+    
     net.save_graph(file_path)
+    with open(file_path, 'r', encoding='utf-8') as f:
+        content = f.read()
+    # Insert the legend before the closing body tag
+    content = content.replace('</body>', f'{legend_html}</body>')
+    with open(file_path, 'w', encoding='utf-8') as f:
+        f.write(content)
 
     return file_path, visualized_edges