diff --git a/Dockerfile b/Dockerfile index beeee59b86cd7a59938fbb49d5e89052f04b3a67..0df07beeabeac47ec67bb4bbb0655668ecef1b3f 100644 --- a/Dockerfile +++ b/Dockerfile @@ -2,9 +2,10 @@ FROM ubuntu:jammy RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get -yq install python3 COPY simulator.py /simulator/ COPY simulator_test.py /simulator/ +COPY dt_model.joblib /model/ WORKDIR /simulator RUN ./simulator_test.py COPY messages.mllp /data/ EXPOSE 8440 EXPOSE 8441 -CMD /simulator/simulator.py --messages=/data/messages.mllp \ No newline at end of file +CMD /simulator/simulator.py --messages=/data/messages.mllp diff --git a/constants.py b/constants.py index d79e0057af20f5d32eac118323fbcc9227741ed4..2bfb8fb4e8e39b27bd5c39707b3566e22225601f 100644 --- a/constants.py +++ b/constants.py @@ -1,2 +1,12 @@ +# MLLP constants MLLP_START_CHAR = b"\x0b" MLLP_END_CHAR = b"\x1c\x0d" + +# Path to load and store the trained Decision Tree model +DT_MODEL_PATH = "model/dt_model.joblib" + +# Map for AKI Label +LABELS_MAP = {"n": 0, "y": 1} + +# Reverse labels map for writing the final output +REVERSE_LABELS_MAP = {v: k for k, v in LABELS_MAP.items()} diff --git a/dt_model.joblib b/dt_model.joblib new file mode 100644 index 0000000000000000000000000000000000000000..625b848ec5c7afd1b9eb8ce0f08aef3135ecc774 Binary files /dev/null and b/dt_model.joblib differ diff --git a/main.py b/main.py index 8ebcbbedffc1f9582e225aa7401d4635c125bea8..a0fd2d09a453081416456924759e635c90ac68f3 100644 --- a/main.py +++ b/main.py @@ -1,12 +1,19 @@ import socket - +from joblib import load from utils import process_mllp_message, parse_hl7_message, create_acknowledgement +from constants import DT_MODEL_PATH, REVERSE_LABELS_MAP + def start_server(host="0.0.0.0", port=8440): """ Starts the TCP server to listen for incoming MLLP messages on the specified port. """ + # Load the model once for use through out + dt_model = load(DT_MODEL_PATH) + assert dt_model != None, "Model is not loaded properly..." + + # Start the server with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: sock.connect((host, port)) print(f"Connected to simulator on {host}:{port}") diff --git a/utils.py b/utils.py index 3503c338ed04af7de065b8887ed8fd87346215c1..69dccf52d9c6f399547b2928201f42be18274346 100644 --- a/utils.py +++ b/utils.py @@ -2,7 +2,7 @@ import socket import hl7 import datetime -from constants import MLLP_START_CHAR, MLLP_END_CHAR +from constants import MLLP_START_CHAR, MLLP_END_CHAR, REVERSE_LABELS_MAP def process_mllp_message(data): @@ -24,7 +24,7 @@ def parse_hl7_message(hl7_data): return message -def create_acknowledgement(hl7_msg): +def create_acknowledgement(): """ Creates an HL7 ACK message for the received message. """ @@ -33,3 +33,27 @@ def create_acknowledgement(hl7_msg): framed_ack = MLLP_START_CHAR + ack_msg.encode() + MLLP_END_CHAR return framed_ack + +def predict_with_dt(dt_model, data): + """ + Following data needs to be passed: + [ + "age", + "sex", + "C1", + "RV1", + "RV1_ratio", + "RV2", + "RV2_ratio", + "change_within_48hrs", + "D" + ] + Predict with the DT Model on the data. + Returns the predicted labels. + """ + y_pred = dt_model.predict(data) + + # Map the predictions to labels + labels = [REVERSE_LABELS_MAP[item] for item in y_pred] + + return labels