DataPrepare/signal_method/shhs_tools.py

62 lines
2.1 KiB
Python

import xml.etree.ElementTree as ET
ANNOTATION_MAP = {
"Wake|0": 0,
"Stage 1 sleep|1": 1,
"Stage 2 sleep|2": 2,
"Stage 3 sleep|3": 3,
"Stage 4 sleep|4": 4,
"REM sleep|5": 5,
"Unscored|9": 9,
"Movement|6": 6
}
SA_EVENTS = ['Central apnea', 'Hypopnea', 'Obstructive apnea']
def parse_sleep_annotations(annotation_path):
"""解析睡眠分期注释"""
try:
tree = ET.parse(annotation_path)
root = tree.getroot()
events = []
for scored_event in root.findall('.//ScoredEvent'):
event_type = scored_event.find('EventType').text
if event_type != "Stages|Stages":
continue
description = scored_event.find('EventConcept').text
start = float(scored_event.find('Start').text)
duration = float(scored_event.find('Duration').text)
if description not in ANNOTATION_MAP:
continue
events.append({
'onset': start,
'duration': duration,
'description': description,
'stage': ANNOTATION_MAP[description]
})
return events
except Exception as e:
return None
def extract_osa_events(annotation_path):
"""提取睡眠呼吸暂停事件"""
try:
tree = ET.parse(annotation_path)
root = tree.getroot()
events = []
for scored_event in root.findall('.//ScoredEvent'):
event_concept = scored_event.find('EventConcept').text
event_type = event_concept.split('|')[0].strip()
if event_type in SA_EVENTS:
start = float(scored_event.find('Start').text)
duration = float(scored_event.find('Duration').text)
if duration >= 10:
events.append({
'start': start,
'duration': duration,
'end': start + duration,
'type': event_type
})
return events
except Exception as e:
return []