app.py

from pathlib import Path
import numpy as np
import gradio as gr
import requests
import json
from transformers import ViTImageProcessor, ViTModel
from PIL import Image

# Store the server's URL
SERVER_URL = "https://affordable-prot-bind-clarke.trycloudflare.com/"
CURRENT_DIR = Path(__file__).parent
DEPLOYMENT_DIR = CURRENT_DIR / "deployment_files"
KEYS_DIR = DEPLOYMENT_DIR / ".fhe_keys"
CLIENT_DIR = DEPLOYMENT_DIR / "client_dir"
SERVER_DIR = DEPLOYMENT_DIR / "server_dir"


USER_ID = "user_id"
EXAMPLE_CLINICAL_TRIAL_LINK = "https://www.trials4us.co.uk/ongoing-clinical-trials/recruiting-healthy-adults-c23026?_gl=1*1ysp815*_up*MQ..&gclid=Cj0KCQjwr9m3BhDHARIsANut04bHqi5zE3sjS3f8JK2WRN3YEgY4bTfWbvTdZTxkUTSISxXX5ZWL7qEaAowwEALw_wcB&gbraid=0AAAAAD3Qci2k_3IERmM6U1FGDuYVayZWH"




# Define possible categories for fields without predefined categories
additional_categories = {
    "Gender": ["Male", "Female", "Other"],
    "Ethnicity": ["White", "Black or African American", "Asian", "American Indian or Alaska Native", "Native Hawaiian or Other Pacific Islander", "Other"],
    "Geographic_Location": ["North America", "South America", "Europe", "Asia", "Africa", "Australia", "Antarctica"],
    "Smoking_Status": ["Never", "Former", "Current"],
    "Diagnoses_ICD10": ["Actinic keratosis", "Melanoma", "Dermatofibroma", "Vascular lesion","None"],
    "Medications": ["Metformin", "Lisinopril", "Atorvastatin", "Amlodipine", "Omeprazole", "Simvastatin", "Levothyroxine", "None"],
    "Allergies": ["Penicillin", "Peanuts", "Shellfish", "Latex", "Bee stings", "None"],
    "Previous_Treatments": ["Chemotherapy", "Radiation Therapy", "Surgery", "Physical Therapy", "Immunotherapy", "None"],
    "Alcohol_Consumption": ["None", "Occasionally", "Regularly", "Heavy"],
    "Exercise_Habits": ["Sedentary", "Light", "Moderate", "Active", "Very Active"],
    "Diet": ["Omnivore", "Vegetarian", "Vegan", "Pescatarian", "Keto", "Mediterranean"],
    "Functional_Status": ["Independent", "Assisted", "Dependent"],
    "Previous_Trial_Participation": ["Yes", "No"]
}

# Define the input components for the form
age_input = gr.Slider(minimum=18, maximum=100, label="Age ", step=1, value=30)
gender_input = gr.Radio(choices=additional_categories["Gender"], label="Gender", value="Male")
ethnicity_input = gr.Radio(choices=additional_categories["Ethnicity"], label="Ethnicity", value="White")
geographic_location_input = gr.Radio(choices=additional_categories["Geographic_Location"], label="Geographic Location", value="North America")
medications_input = gr.CheckboxGroup(choices=additional_categories["Medications"], label="Medications", value=["Metformin"])
allergies_input = gr.CheckboxGroup(choices=additional_categories["Allergies"], label="Allergies", value=["Peanuts"])
previous_treatments_input = gr.CheckboxGroup(choices=additional_categories["Previous_Treatments"], label="Previous Treatments", value=["None"])
blood_glucose_level_input = gr.Slider(minimum=0, maximum=300, label="Blood Glucose Level", step=1, value=100)
blood_pressure_systolic_input = gr.Slider(minimum=80, maximum=200, label="Blood Pressure (Systolic)", step=1, value=120)
blood_pressure_diastolic_input = gr.Slider(minimum=40, maximum=120, label="Blood Pressure (Diastolic)", step=1, value=80)
bmi_input = gr.Slider(minimum=10, maximum=50, label="BMI ", step=1, value=20)
smoking_status_input = gr.Radio(choices=additional_categories["Smoking_Status"], label="Smoking Status", value="Never")
alcohol_consumption_input = gr.Radio(choices=additional_categories["Alcohol_Consumption"], label="Alcohol Consumption", value="None")
exercise_habits_input = gr.Radio(choices=additional_categories["Exercise_Habits"], label="Exercise Habits", value="Sedentary")
diet_input = gr.Radio(choices=additional_categories["Diet"], label="Diet", value="Omnivore")
condition_severity_input = gr.Slider(minimum=1, maximum=10, label="Condition Severity", step=1, value=5)
functional_status_input = gr.Radio(choices=additional_categories["Functional_Status"], label="Functional Status", value="Independent")
previous_trial_participation_input = gr.Radio(choices=additional_categories["Previous_Trial_Participation"], label="Previous Trial Participation", value="Yes")


# def encrypt_array(user_symptoms: np.ndarray, user_id: str) -> bytes:
#     """
#     Encrypt the user symptoms vector.

#     Args:
#         user_symptoms (np.ndarray): The vector of symptoms provided by the user.
#         user_id (str): The current user's ID.
    
#     Returns:
#         bytes: Encrypted and serialized symptoms.
#     """

#     # Retrieve the client API
#     client = FHEModelClient(path_dir=DEPLOYMENT_DIR, key_dir=KEYS_DIR / f"{user_id}")
#     client.load()

#     # Ensure the symptoms are properly formatted as an array
#     user_symptoms = np.array(user_symptoms).reshape(1, -1)

#     # Encrypt and serialize the symptoms
#     encrypted_quantized_user_symptoms = client.quantize_encrypt_serialize(user_symptoms)
    
#     # Ensure the encryption process returned bytes
#     assert isinstance(encrypted_quantized_user_symptoms, bytes)

#     # Save the encrypted data to a file (optional)
#     encrypted_input_path = KEYS_DIR / f"{user_id}/encrypted_input"
#     with encrypted_input_path.open("wb") as f:
#         f.write(encrypted_quantized_user_symptoms)

#     # Return the encrypted data
#     return encrypted_quantized_user_symptoms


# def decrypt_result(encrypted_answer: bytes, user_id: str) -> bool:
"""
    Decrypt the encrypted result.

    Args:
        encrypted_answer (bytes): The encrypted result.
        user_id (str): The current user's ID.
    
    Returns:
        bool: The decrypted result.
    """

    # Retrieve the client API
    # client = FHEModelClient(path_dir=DEPLOYMENT_DIR, key_dir=KEYS_DIR / f"{user_id}")
    # client.load()

    # Decrypt the result
    # decrypted_result = client.decrypt_deserialize(encrypted_answer)

    # # Return the decrypted result
    # return decrypted_result



def encode_categorical_data(data):
    categories = ["Gender", "Ethnicity", "Geographic_Location", "Diagnoses_ICD10", "Medications", "Allergies", "Previous_Treatments", "Smoking_Status", "Alcohol_Consumption", "Exercise_Habits", "Diet", "Functional_Status", "Previous_Trial_Participation"]
    encoded_data = []
    for i in range(len(categories)):
        sub_cats = additional_categories[categories[i]]
        if data[i] in sub_cats:
            encoded_data.append(sub_cats.index(data[i]) + 1)
        else:
            encoded_data.append(0)
        
    return encoded_data

def clear_data_to_json(data):
    print(data)
    patient_data = {
        "model_names": ["second_model"],
        "patient": {
            "Age": data.get("age", 30),
            "Blood_Glucose_Level": data.get("blood_glucose_level", 0),
            "Blood_Pressure_Systolic": data.get("blood_pressure_systolic", 0),
            "Blood_Pressure_Diastolic": data.get("blood_pressure_diastolic", 0),
            "BMI": data.get("bmi", 0),
            "Condition_Severity": data.get("condition_severity", 0),
            "Gender": data.get("Gender", 0),
            "Ethnicity": data.get("Ethnicity", 0),
            "Geographic_Location": data.get("Geographic_Location", 0),
            "Smoking_Status": data.get("Smoking_Status", 0),
            "Diagnoses_ICD10": data.get("Diagnoses_ICD10", 0),
            "Medications": data.get("Medications", 0),
            "Allergies": data.get("Allergies", 0),
            "Previous_Treatments": data.get("Previous_Treatments", 0),
            "Alcohol_Consumption": data.get("Alcohol_Consumption", 0),
            "Exercise_Habits": data.get("Exercise_Habits", 0),
            "Diet": data.get("Diet", 0),
            "Functional_Status": data.get("Functional_Status", 0),
            "Previous_Trial_Participation": data.get("Previous_Trial_Participation", 0)
        }
    }
    return patient_data


def process_patient_data(age, gender, ethnicity, geographic_location, diagnoses_icd10, medications, allergies, previous_treatments, blood_glucose_level, blood_pressure_systolic, blood_pressure_diastolic, bmi, smoking_status, alcohol_consumption, exercise_habits, diet, condition_severity, functional_status, previous_trial_participation):    
    
    # Encode the data
    categorical_data = [gender, ethnicity, geographic_location, diagnoses_icd10, medications, allergies, previous_treatments, smoking_status, alcohol_consumption, exercise_habits, diet, functional_status, previous_trial_participation]
    print(f"Categorical data: {categorical_data}")
    encoded_categorical_data = encode_categorical_data(categorical_data)
    numerical_data = np.array([age, blood_glucose_level, blood_pressure_systolic, blood_pressure_diastolic, bmi, condition_severity])
    print(f"Numerical data: {numerical_data}")
    print(f"One-hot encoded data: {encoded_categorical_data}")
    combined_data = np.hstack((numerical_data, encoded_categorical_data))

    ordered_categories = ["Gender", "Ethnicity", "Geographic_Location", "Diagnoses_ICD10", "Medications", "Allergies", "Previous_Treatments", "Smoking_Status", "Alcohol_Consumption", "Exercise_Habits", "Diet", "Functional_Status", "Previous_Trial_Participation"]
    zipped_data = zip(ordered_categories, encoded_categorical_data)
    # Convert the zipped data to a dictionary
    encoded_categorical_dict = {category: value for category, value in zipped_data}

    # Convert the data to JSON
    json_data = clear_data_to_json({
        "age": age,
        "blood_glucose_level": blood_glucose_level,
        "blood_pressure_systolic": blood_pressure_systolic,
        "blood_pressure_diastolic": blood_pressure_diastolic,
        "bmi": bmi,
        "condition_severity": condition_severity,
        **encoded_categorical_dict
    })

    print(f"JSON data: {json_data}")
    print(f"Combined data: {combined_data}")
    # encrypted_array = encrypt_array(combined_data, "user_id")

    # Send the data to the server 
    url = SERVER_URL + "inference/clear-match" 
    headers = {"Content-Type": "application/json", "X-API-KEY": "secret"}
    response = requests.post(url, data=json_data, headers=headers)

    # Check if the data was sent successfully
    if response.status_code == 200:
        print("Data sent successfully.")
    else:
        print(f"Error sending data. Status code: {response.status_code}")
    
    # return f"**[Possible Trial Link]({EXAMPLE_CLINICAL_TRIAL_LINK})**"
    # Decrypt the result
    # decrypted_result = decrypt_result(response.content, USER_ID)
    # decrypted_result = None
    try:
        decrypted_result = response.json()
        print(f"Decrypted result: {decrypted_result}")
    except json.JSONDecodeError as e:
        print(f"Error decrypting result: Invalid JSON. {e}")
        decrypted_result = False
        return "There was an error processing the data."
    except Exception as e:
        print(f"An error occurred: {e}")
        decrypted_result = False
        return "There was an error processing the data."
    else:
        print("Json parsed successfully.")
    
    
     # If the answer is True, return the link
    if decrypted_result:
        return f"**[Possible Trial Link]({EXAMPLE_CLINICAL_TRIAL_LINK})**"
    else:
        return "There was an error processing the data."
    

# Define the function to handle image upload
def handle_image_upload(image):
    url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
    image = Image.open(requests.get(url, stream=True).raw)

    processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224-in21k')
    model = ViTModel.from_pretrained('google/vit-base-patch16-224-in21k')
    inputs = processor(images=image, return_tensors="pt")

    outputs = model(**inputs)
    pooler_output = outputs.pooler_output[0]
    sclaed_output = 127 + 127 * pooler_output / pooler_output.abs().max()
    sclaed_output = sclaed_output.to(int)
    url = "/inference/clear-diagnosis"
    return ["Melanoma", "Vascular lesion"]

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Patient Data Criteria Form\nPlease fill in the criteria for the type of patients you are looking for.")
    with gr.Column():
        with gr.Group():
            age_input.render()
            gender_input.render()
            ethnicity_input.render()
            geographic_location_input.render()
            medications_input.render()
            allergies_input.render()
            previous_treatments_input.render()
            blood_glucose_level_input.render()
            blood_pressure_systolic_input.render()
            blood_pressure_diastolic_input.render()
            bmi_input.render()
            smoking_status_input.render()
            alcohol_consumption_input.render()
            exercise_habits_input.render()
            diet_input.render()
            condition_severity_input.render()
            functional_status_input.render()
            previous_trial_participation_input.render()
        with gr.Group():
            diagnoses_icd10_input = gr.CheckboxGroup(choices=additional_categories["Diagnoses_ICD10"], label="Skin Diagnosis", interactive=False)
            image_input = gr.Image(label="Upload an Image")
            gr.Button("Upload").click(handle_image_upload, inputs=image_input, outputs=diagnoses_icd10_input)
        with gr.Group():
            output = gr.Markdown("**Server response**")
            gr.Button("Submit").click(process_patient_data, inputs=[
                age_input, gender_input, ethnicity_input, geographic_location_input, diagnoses_icd10_input, medications_input, allergies_input, previous_treatments_input, blood_glucose_level_input, blood_pressure_systolic_input, blood_pressure_diastolic_input, bmi_input, smoking_status_input, alcohol_consumption_input, exercise_habits_input, diet_input, condition_severity_input, functional_status_input, previous_trial_participation_input
            ], outputs=output)


# Launch the app
demo.launch()