refactor

README.md

@@ -11,38 +11,115 @@ pinned: false
 
 # Medical Knowledge Graph Construction (medKGC)
 
-## English Version
+## Overview
+medKGC is a medical text knowledge graph construction and review system. It supports entity recognition, relation extraction, and visualization of medical reports, providing a convenient review interface.
 
-medKGC is a Streamlit-based application for medical text analysis and knowledge graph construction. It demonstrates the process of entity recognition and relation extraction from medical texts, visualizing the results in an interactive graph.
+## Core Features
 
-### Features:
-1. Text Input: Users can input medical text for analysis.
-2. Entity Recognition: The app identifies and highlights various medical entities such as diseases, procedures, anatomy, etc.
-3. Relation Extraction: It extracts relationships between the identified entities.
-4. Interactive Visualization: Utilizes streamlit-agraph to create an interactive graph representation of entities and their relationships.
-5. Labeled Text Display: Shows the input text with highlighted entities using streamlit_text_label.
+### 1. Data Processing
+- **Position Conversion**: Support word-level and char-level position conversion
+- **Entity Conversion**: Convert between JSON format and Selection objects
+- **Relation Extraction**: Entity ID-based relation mapping and reconstruction
 
-### How to Use:
-1. Enter medical text in the provided text area.
-2. Click "Recognize Entities" to process the text.
-3. View the recognized entities, extracted relations, and the entity relationship graph.
+### 2. Entity Annotation
+- **Label Types**:
+  - OBS-DP: Observation definitely present (Red)
+  - ANAT-DP: Anatomy definitely present (Cyan)
+  - OBS-U: Observation uncertain (Yellow)
+  - OBS-DA: Observation definitely absent (Gray)
+- **Interactive Annotation**: Support entity selection and annotation
 
-Note: Currently, the app uses mock data for demonstration purposes. Integration with actual NLP models is planned for future development.
+### 3. Relation Visualization
+- **Node Merging**: Automatically merge entities with same text
+- **Color Coding**: Different colors for different entity types
+- **Dynamic Updates**: Support real-time graph updates
 
-## 中文版本
+### 4. Review Process
+- **Report Selection**: Display pending and reviewed reports separately
+- **Status Saving**: Automatically save review status and modifications
+- **Batch Processing**: Support continuous review of multiple reports
 
-medKGC 是一个基于 Streamlit 的医疗文本分析和知识图谱构建应用。它演示了从医疗文本中进行实体识别和关系提取的过程，并将结果以交互式图形可视化。
+## Technical Implementation
 
-### 功能特点：
-1. 文本输入：用户可以输入医疗文本进行分析。
-2. 实体识别：应用程序识别并高亮显示各种医疗实体，如疾病、医疗程序、解剖结构等。
-3. 关系提取：提取识别出的实体之间的关系。
-4. 交互式可视化：使用 streamlit-agraph 创建实体及其关系的交互式图形表示。
-5. 标记文本显示：使用 streamlit_text_label 显示带有高亮实体的输入文本。
+### Data Structures
+1. **Entity Data**
+```json
+{
+    "entities": {
+        "1": {
+            "tokens": "entity text",
+            "label": "entity type",
+            "start_ix": "word-level start position",
+            "end_ix": "word-level end position",
+            "relations": [["relation type", "target entity ID"]]
+        }
+    }
+}
+```
 
-### 使用方法：
-1. 在提供的文本区域输入医疗文本。
-2. 点击"识别实体"按钮处理文本。
-3. 查看识别出的实体、提取的关系以及实体关系图。
+2. **Selection Object**
+```python
+@dataclass
+class Selection:
+    start: int  # char-level start position
+    end: int    # char-level end position
+    text: str   # entity text
+    labels: List[str]  # entity type list
+```
 
-注意：目前，应用程序使用模拟数据进行演示。未来开发计划将集成实际的自然语言处理模型。
+### Core Algorithms
+1. **Position Conversion**
+```python
+def word_to_char_span(text, start_ix, end_ix):
+    """Convert word-level position to character-level range"""
+```
+
+2. **Relation Reconstruction**
+```python
+def find_relations_with_entities(entities, entities_data):
+    """Rebuild relations based on entity text matching"""
+```
+
+## Deployment
+
+### Requirements
+- Python 3.7+
+- Streamlit 1.39.0+
+- streamlit_text_label
+- streamlit_agraph
+
+### Installation
+1. Clone repository
+```bash
+git clone https://github.com/your-repo/medKGC.git
+```
+
+2. Install dependencies
+```bash
+pip install -r requirements.txt
+```
+
+3. Run application
+```bash
+streamlit run app.py
+```
+
+## Future Plans
+1. [ ] Add relation editing functionality
+2. [ ] Support custom entity types
+3. [ ] Add data export functionality
+4. [ ] Integrate machine learning models
+5. [ ] Add annotation functionality
+
+## Contributing
+Welcome to contribute through:
+1. Submit Issues for bug reports or suggestions
+2. Submit Pull Requests to improve code
+3. Improve documentation and comments
+
+## License
+MIT License
+
+---
+
+[Chinese version above]

app.py

@@ -1,91 +1,275 @@
 import streamlit as st
+import json
 from streamlit_text_label import label_select, Selection
 from streamlit_agraph import agraph, Node, Edge, Config
-
-# Import Relation class
 from dataclasses import dataclass
 
+
 @dataclass
 class Relation:
     source: Selection
     target: Selection
     label: str
 
-def mock_entity_recognition(text):
-    # Simulate entity recognition functionality
-    entities = [
-        Selection(start=0, end=12, text="FINAL REPORT",
-                  labels=["REPORT_TYPE"]),
-        Selection(start=22, end=33, text="Pseudomonas", labels=["DISEASE"]),
-        Selection(start=39, end=49, text="intubation", labels=["PROCEDURE"]),
-        Selection(start=62, end=67, text="study", labels=["EXAM"]),
-        Selection(start=116, end=126, text="Substantial", labels=["SEVERITY"]),
-        Selection(start=127, end=145, text="bilateral pleural",
-                  labels=["ANATOMY"]),
-        Selection(start=146, end=155, text="effusions", labels=["OBSERVATION"])
-    ]
-    return entities
 
+def load_data():
+    """Load data from dev.json"""
+    with open('dev.json', 'r') as f:
+        return json.load(f)
 
-def mock_relation_extraction(entities):
-    # Simulate relation extraction functionality
-    relations = [
-        Relation(source=entities[0],
-                 target=entities[1], label="DISEASE_CAUSE"),
-        Relation(source=entities[1], target=entities[2],
-                 label="PROCEDURE_EFFECT"),
-        Relation(source=entities[2], target=entities[3], label="EXAM_RESULT"),
-        Relation(source=entities[3], target=entities[4],
-                 label="SEVERITY_LEVEL"),
-        Relation(source=entities[4], target=entities[5],
-                 label="ANATOMY_LOCATION"),
-    ]
 
-    return relations
+def save_data(data):
+    """Save data to dev.json"""
+    with open('dev.json', 'w') as f:
+        json.dump(data, f, indent=4)
+
+
+def get_label_color(label):
+    """Return color based on label type"""
+    color_map = {
+        'OBS-DP': '#FF6B6B',  # Red - Observation definitely present
+        'ANAT-DP': '#4ECDC4',  # Cyan - Anatomy definitely present
+        'OBS-U': '#FFD93D',    # Yellow - Observation uncertain
+        'OBS-DA': '#95A5A6',   # Gray - Observation definitely absent
+    }
+    return color_map.get(label, '#666666')  # Default color
 
 
 def create_graph(entities, relations):
-    nodes = [Node(id=e.text, label=e.text, size=25, color=f"#{hash(e.labels[0]) % 0xFFFFFF:06x}") for e in entities]
-    edges = [Edge(source=r.source.text, target=r.target.text, label=r.label) for r in relations]
-    
-    config = Config(width=750, height=500, directed=True, physics=True, hierarchical=False)
+    """Create entity relationship graph, merge nodes with same text"""
+    # Track created nodes using dict, key is entity text
+    nodes_dict = {}
+    nodes = []
+
+    # First create all unique nodes
+    for entity in entities:
+        if entity.text not in nodes_dict:
+            # Create new node
+            node = Node(
+                id=entity.text,
+                label=f"{entity.text}\n({entity.labels[0]})",
+                size=25,
+                color=get_label_color(entity.labels[0])
+            )
+            nodes.append(node)
+            nodes_dict[entity.text] = node
+
+    # Create edges using node text as source and target
+    edges = []
+    for relation in relations:
+        # Check if source and target exist
+        if relation.source.text in nodes_dict and relation.target.text in nodes_dict:
+            edge = Edge(
+                source=relation.source.text,
+                target=relation.target.text,
+                label=relation.label,
+                color="#666666"  # Unified edge color
+            )
+            edges.append(edge)
+
+    config = Config(
+        width=750,
+        height=500,
+        directed=True,
+        physics=True,
+        hierarchical=False,
+        nodeHighlightBehavior=True,
+        highlightColor="#F7A7A6",
+    )
+
     return agraph(nodes=nodes, edges=edges, config=config)
 
 
+def word_to_char_position(text, word_index):
+    """Convert word position to character position"""
+    words = text.split()
+    char_start = 0
+
+    # If word_index out of range, return text end
+    if word_index >= len(words):
+        return len(text)
+
+    # Traverse all words before target word
+    for i in range(word_index):
+        char_start += len(words[i]) + 1  # +1 for space
+
+    return char_start
+
+
+def word_to_char_span(text, start_ix, end_ix):
+    """Convert word start and end positions to character span"""
+    char_start = word_to_char_position(text, start_ix)
+    # If start equals end, it's a single word
+    if start_ix == end_ix:
+        char_end = char_start + len(text.split()[start_ix])
+    else:
+        # If multiple words, calculate to end position
+        char_end = word_to_char_position(
+            text, end_ix) + len(text.split()[end_ix])
+    return char_start, char_end
+
+
+def entities2Selection(text, entities_data):
+    """Convert entities data to Selection objects list"""
+    selections = []
+    for entity_id, entity in entities_data.items():
+        # Convert word positions to char positions
+        char_start, char_end = word_to_char_span(
+            text,
+            entity['start_ix'],
+            entity['end_ix']
+        )
+
+        selection = Selection(
+            start=char_start,
+            end=char_end,
+            text=entity['tokens'],
+            labels=[entity['label']],
+        )
+        selections.append(selection)
+    return selections
+
+
+def selection2entities(selections):
+    """Convert Selection objects list to entities data"""
+    entities = {}
+    for i, selection in enumerate(selections, 1):
+        entities[str(i)] = {
+            "tokens": selection.text,
+            "label": selection.labels[0],
+            "start_ix": selection.start,
+            "end_ix": selection.end,
+            "relations": []  # Initialize empty relations list
+        }
+    return entities
+
+
+def find_relations_with_entities(entities, entities_data):
+    """Find relations between current entities based on original entities_data"""
+    # Create text to entity mapping
+    text_to_entity = {e.text: e for e in entities}
+
+    # Create tokens to entity_id mapping
+    tokens_to_id = {entity['tokens']: entity_id
+                    for entity_id, entity in entities_data.items()}
+
+    # Create id to tokens mapping
+    id_to_tokens = {entity_id: entity['tokens']
+                    for entity_id, entity in entities_data.items()}
+
+    relations = []
+    # Iterate through each entity in current entities
+    for source_text, source_entity in text_to_entity.items():
+        # Find corresponding entity ID in original data
+        for entity_id, entity in entities_data.items():
+            if entity['tokens'] == source_text:
+                # Iterate through all relations of this entity
+                for relation in entity.get('relations', []):
+                    target_id = relation[1]
+                    # Get target entity text
+                    target_text = id_to_tokens.get(target_id)
+                    # If target entity exists in current entities
+                    if target_text and target_text in text_to_entity:
+                        relations.append(Relation(
+                            source=source_entity,
+                            target=text_to_entity[target_text],
+                            label=relation[0]
+                        ))
+
+    return relations
+
+
 def main():
-    st.title("Medical Text Entity Recognition")
-
-    # 1. Text input
-    text = st.text_area("Enter medical text:", value="FINAL REPORT HISTORY : Pseudomonas with intubation . FINDINGS : In comparison with the study of ___ , there is little change in the monitoring and support devices . Substantial bilateral pleural effusions , more prominent on the right with bibasilar atelectasis .")
-
-    if st.button("Recognize Entities"):
-        # 2. Call the simulated entity recognition function
-        entities = mock_entity_recognition(text)
-
-        # 3. Use streamlit_text_label to highlight recognized entities
-        st.subheader("Recognition Results:")
-        labeled_text = label_select(
-            body=text,
-            labels=["REPORT_TYPE", "DISEASE", "PROCEDURE",
-                    "EXAM", "SEVERITY", "ANATOMY", "OBSERVATION"],
-            selections=entities
+    """Main application"""
+    st.title("Medical Report Review System")
+
+    # Load data
+    if 'data' not in st.session_state:
+        st.session_state.data = load_data()
+
+    # Create two columns layout
+    col1, col2 = st.columns(2)
+
+    with col1:
+        st.subheader("Reports to Review")
+        # Get unreviewed reports list
+        unreviewed_reports = [
+            report_id for report_id, content in st.session_state.data.items()
+            if 'reviewed' not in content
+        ]
+        selected_report = st.selectbox(
+            "Select Report",
+            unreviewed_reports,
+            key="unreviewed"
         )
 
-        st.write("Recognized entities:")
-        for entity in entities:
-            st.write(f"{entity.text} ({entity.labels[0]})")
+    with col2:
+        st.subheader("Reviewed Reports")
+        # Get reviewed reports list
+        reviewed_reports = [
+            report_id for report_id, content in st.session_state.data.items()
+            if content.get('reviewed', False)
+        ]
+        st.selectbox(
+            "Completed Reports",
+            reviewed_reports if reviewed_reports else ['None'],
+            key="reviewed"
+        )
+
+    if selected_report:
+        report_data = st.session_state.data[selected_report]
+
+        # Display report text
+        st.subheader("Report Content:")
+        st.markdown(report_data['text'])
 
-        # 4. Call the simulated relation extraction function
-        relations = mock_relation_extraction(entities)
+        # Display entities and relations
+        entities_data = report_data['entities']
 
-        # 5. Display relations
-        st.subheader("Extracted Relations:")
-        for relation in relations:
-            st.write(f"{relation.source.text} --{relation.label}--> {relation.target.text}")
+        # Convert entities data to Selection objects
+        entities = entities2Selection(report_data['text'], entities_data)
 
-        # 6. Create and display graph using streamlit-agraph
+        # Display annotation results
+        st.subheader("Entity Annotation:")
+        selections = label_select(
+            body=report_data['text'],
+            labels=list(set(e.labels[0] for e in entities)),
+            selections=entities,
+        )
+        # Display selections in text format
+        st.write(selections)
+
+        # Display relationship graph
         st.subheader("Entity Relationship Graph:")
-        create_graph(entities, relations)
+
+        # Add update graph button
+        if st.button("Update Graph"):
+            # Update relations using current selections
+            relations = find_relations_with_entities(selections, entities_data)
+            create_graph(selections, relations)
+        else:
+            # Display original graph
+            relations = find_relations_with_entities(entities, entities_data)
+            create_graph(entities, relations)
+
+        # Add review functionality
+        if st.button("Mark as Reviewed"):
+            # Convert selections back to entities data
+            updated_entities = selection2entities(selections)
+
+            # Keep original relations
+            for entity_id, entity in updated_entities.items():
+                if entity_id in entities_data:
+                    entity['relations'] = entities_data[entity_id]['relations']
+
+            st.session_state.data[selected_report]['reviewed'] = {
+                'entities': updated_entities
+            }
+
+            save_data(st.session_state.data)
+            st.success("Review status saved!")
+            st.rerun()
+
 
 if __name__ == "__main__":
     main()

mockData.json

@@ -1,138 +0,0 @@
-{'data_source': 'MIMIC-CXR',
- 'data_split': 'train',
- 'entities': {'1': {'end_ix': 29,
-                    'label': 'OBS-DP',
-                    'relations': [['modify', '4']],
-                    'start_ix': 29,
-                    'tokens': 'Substantial'},
-              '10': {'end_ix': 44,
-                     'label': 'OBS-DP',
-                     'relations': [['suggestive_of', '16']],
-                     'start_ix': 44,
-                     'tokens': 'configuration'},
-              '11': {'end_ix': 47,
-                     'label': 'OBS-DP',
-                     'relations': [['modify', '12']],
-                     'start_ix': 47,
-                     'tokens': 'collection'},
-              '12': {'end_ix': 49,
-                     'label': 'OBS-DP',
-                     'relations': [['located_at', '14'],
-                                   ['suggestive_of', '16']],
-                     'start_ix': 49,
-                     'tokens': 'opacification'},
-              '13': {'end_ix': 52,
-                     'label': 'ANAT-DP',
-                     'relations': [['modify', '14']],
-                     'start_ix': 52,
-                     'tokens': 'left'},
-              '14': {'end_ix': 53,
-                     'label': 'ANAT-DP',
-                     'relations': [],
-                     'start_ix': 53,
-                     'tokens': 'base'},
-              '15': {'end_ix': 59,
-                     'label': 'OBS-U',
-                     'relations': [['modify', '16']],
-                     'start_ix': 59,
-                     'tokens': 'loculated'},
-              '16': {'end_ix': 60,
-                     'label': 'OBS-U',
-                     'relations': [],
-                     'start_ix': 60,
-                     'tokens': 'fluid'},
-              '17': {'end_ix': 67,
-                     'label': 'OBS-DP',
-                     'relations': [['modify', '20']],
-                     'start_ix': 67,
-                     'tokens': 'increased'},
-              '18': {'end_ix': 68,
-                     'label': 'ANAT-DP',
-                     'relations': [['modify', '19']],
-                     'start_ix': 68,
-                     'tokens': 'pulmonary'},
-              '19': {'end_ix': 69,
-                     'label': 'ANAT-DP',
-                     'relations': [],
-                     'start_ix': 69,
-                     'tokens': 'venous'},
-              '2': {'end_ix': 30,
-                    'label': 'ANAT-DP',
-                    'relations': [['modify', '3']],
-                    'start_ix': 30,
-                    'tokens': 'bilateral'},
-              '20': {'end_ix': 70,
-                     'label': 'OBS-DP',
-                     'relations': [['located_at', '19']],
-                     'start_ix': 70,
-                     'tokens': 'pressure'},
-              '21': {'end_ix': 79,
-                     'label': 'ANAT-DP',
-                     'relations': [['modify', '23']],
-                     'start_ix': 79,
-                     'tokens': 'left'},
-              '22': {'end_ix': 80,
-                     'label': 'ANAT-DP',
-                     'relations': [['modify', '23']],
-                     'start_ix': 80,
-                     'tokens': 'upper'},
-              '23': {'end_ix': 81,
-                     'label': 'ANAT-DP',
-                     'relations': [],
-                     'start_ix': 81,
-                     'tokens': 'zone'},
-              '24': {'end_ix': 87,
-                     'label': 'ANAT-DP',
-                     'relations': [],
-                     'start_ix': 87,
-                     'tokens': 'cavitary'},
-              '25': {'end_ix': 88,
-                     'label': 'OBS-DP',
-                     'relations': [['located_at', '24']],
-                     'start_ix': 88,
-                     'tokens': 'process'},
-              '3': {'end_ix': 31,
-                    'label': 'ANAT-DP',
-                    'relations': [],
-                    'start_ix': 31,
-                    'tokens': 'pleural'},
-              '4': {'end_ix': 32,
-                    'label': 'OBS-DP',
-                    'relations': [['located_at', '3']],
-                    'start_ix': 32,
-                    'tokens': 'effusions'},
-              '5': {'end_ix': 35,
-                    'label': 'OBS-DP',
-                    'relations': [['located_at', '6']],
-                    'start_ix': 35,
-                    'tokens': 'prominent'},
-              '6': {'end_ix': 38,
-                    'label': 'ANAT-DP',
-                    'relations': [['modify', '3']],
-                    'start_ix': 38,
-                    'tokens': 'right'},
-              '7': {'end_ix': 40,
-                    'label': 'ANAT-DP',
-                    'relations': [],
-                    'start_ix': 40,
-                    'tokens': 'bibasilar'},
-              '8': {'end_ix': 41,
-                    'label': 'OBS-DP',
-                    'relations': [['located_at', '6'], ['located_at', '7']],
-                    'start_ix': 41,
-                    'tokens': 'atelectasis'},
-              '9': {'end_ix': 43,
-                    'label': 'OBS-DP',
-                    'relations': [['modify', '10']],
-                    'start_ix': 43,
-                    'tokens': 'Unusual'}},
- 'text': 'FINAL REPORT HISTORY : Pseudomonas with intubation . FINDINGS : In '
-         'comparison with the study of ___ , there is little change in the '
-         'monitoring and support devices . Substantial bilateral pleural '
-         'effusions , more prominent on the right with bibasilar atelectasis . '
-         'Unusual configuration to the collection of opacification at the left '
-         'base raises the possibility of some loculated fluid . There is again '
-         'evidence of increased pulmonary venous pressure . Overlapping '
-         'structures somewhat obscure visualization of the left upper zone and '
-         'simulate the appearance of cavitary process . This area should be '
-         'closely checked on subsequent radiographs .'}