app.py

import streamlit as st
from datetime import datetime
from diffusers import AutoPipelineForText2Image, DPMSolverMultistepScheduler
import matplotlib.pyplot as plt
import torch
import re
import os
import sys
import time

import smtplib
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.mime.application import MIMEApplication
from pydrive.drive import GoogleDrive
from pydrive.auth import GoogleAuth

from PIL import Image, ImageDraw, ImageFont

def add_annotation(image, text, font_path=None, font_size=20, margin=[10,10]):
    # Create a new image for the annotation
    annotation_height = font_size + margin[0] + margin[1]  # Height of the annotation area
    annotation_image = Image.new('RGB', (image.width, annotation_height), 'black')
    
    # Draw the text on the annotation image
    draw = ImageDraw.Draw(annotation_image)
    font = ImageFont.truetype(font_path, font_size) if font_path else ImageFont.load_default()
    _, _, text_width, text_height = draw.textbbox((0, 0), text, font=font)

    text_bbox = draw.textbbox((0, 0), text, font=font)
    text_width, text_height = text_bbox[2] - text_bbox[0], text_bbox[3] - text_bbox[1]
    text_position = ((image.width - text_width) // 2, margin[0] + (annotation_height - margin[0] - margin[1] - text_height) // 2)

    draw.text(text_position, text, fill="white", font=font)

    
    # Create a new image by combining the original image and the annotation
    combined_image = Image.new('RGB', (image.width, image.height + annotation_height))
    combined_image.paste(image, (0, 0))
    combined_image.paste(annotation_image, (0, image.height))
    
    return combined_image

def drive_authenticate(cred_filename = "mycreds.txt"):
    gauth = GoogleAuth()
    gauth.LoadCredentialsFile(cred_filename)
    if gauth.credentials is None:
        # Authenticate if they're not there
        gauth.LocalWebserverAuth()
    elif gauth.access_token_expired:
        # Refresh them if expired
        gauth.Refresh()
    else:
        # Initialize the saved creds
        gauth.Authorize()
    # Save the current credentials to a file
    gauth.SaveCredentialsFile(cred_filename)
    return gauth

def drive_download(drive, folder_name, file_name):
    # Search for the folder
    folder_list = drive.ListFile({'q': f"title='{folder_name}' and mimeType='application/vnd.google-apps.folder' and trashed=false"}).GetList()
    if not folder_list:
        print(f"Folder '{folder_name}' not found.")
        return

    folder_id = folder_list[0]['id']

    # Search for the file within the folder
    file_list = drive.ListFile({'q': f"title='{file_name}' and '{folder_id}' in parents and trashed=false"}).GetList()
    if not file_list:
        print(f"File '{file_name}' not found in folder '{folder_name}'.")
        return

    file = file_list[0]
    file.GetContentFile(file_name)
    print(f'Downloaded file: {file_name}')

def drive_search_folder(drive, folder_name):
    # Search for the folder
    folder_list = drive.ListFile({'q': f"title='{folder_name}' and mimeType='application/vnd.google-apps.folder' and trashed=false"}).GetList()
    if not folder_list:
        print(f"Folder '{folder_name}' not found.")
        return 0

    folder_id = folder_list[0]['id']
    return folder_id

def drive_search_file(drive, folder_id, file_name):
    # Search for the file within the folder
    file_list = drive.ListFile({'q': f"title='{file_name}' and '{folder_id}' in parents and trashed=false"}).GetList()
    if not file_list:
        print(f"File '{file_name}' not found in folder '{folder_id}'.")
        return 0

    file = file_list[0]
    return file

def drive_update(drive, folder_name, file_name):
    folder_id = drive_search_folder(drive, folder_name)
    file = drive_search_file(drive, folder_id, file_name)
    if file:
        file.SetContentFile(file_name)
        file.Upload()
        print(f'File {file_name} updated!')
    else:
        drive_upload(drive, folder_name, file_name)

def drive_upload(drive, folder_name, file_name):
    folder_id = drive_search_folder(drive, folder_name)
    gfile = drive.CreateFile({'parents': [{'id': folder_id}]}) # Read file and set it as the content of this instance. 
    gfile.SetContentFile(file_name)
    gfile.Upload() # Upload the file.
    print(f'File {file_name} uploaded!')

drive_folder_name = 'SIRM2024 Pictures'
drive_txt_name = 'emails.txt'

sender = "deeptracetech.sirm2024@gmail.com"
privacy_statement = '''
    DEEP TRACE TECHNOLOGIES Srl cares about the privacy of its clients and potential clients. Therefore, we ensure that all data you provide by filling out the form or voluntarily providing your contact information will be treated with the utmost confidentiality.\n
    In compliance with the regulations on the protection of personal data, we inform you that the collected data will be processed in a manner that ensures security and confidentiality.
    The personal data you provide will be processed for the following purposes:
    1. Creation of a personalized profile related to the use of this application.
    2. Sending commercial communications regarding activities related to the requested services or sending information about new products.
    3. Communicating your contact details to third parties (such as agents or sales representatives in your area) to meet your requests.\n
    Our collaborators and/or staff may become aware of the data in question as data controllers or processors.\n
    Your rights, as provided by the regulations, can be exercised by submitting a written request to the data controller at the following email address: admin@deeptracetech.com\n
    DeepTrace Technologies S.R.L., Via Conservatorio 17 - 20122 Milan (MI)
'''

mail_body = '''
    Thank you for joining DeepTrace Technologies at SIRM 2024! Attached to this mail, you will find the AI medical image you generated with the style of your favourite artist!\nù
    If you wish not to continue receiving updates on the latest news from DeepTrace Technologies, reply UNSUBSCRIBE to this email.\n
    Best regards, the team of DeepTrace Technologies
'''

if 'button_clicked' not in st.session_state:
    st.session_state.button_clicked = False

# Define a function to handle the button click
def on_button_click():
    st.session_state.button_clicked = True 

def send_email(subject, body, attachments, sender, recipient, **kwargs):

    smtp_server = kwargs.get('smtp_server', 'smtp.gmail.com')

    # MIMEMultipart() creates a container for an email message that can hold
    # different parts, like text and attachments and in next line we are
    # attaching different parts to email container like subject and others.
    message = MIMEMultipart()
    message['Subject'] = subject
    message['From'] = sender
    message['To'] = recipient
    body_part = MIMEText(body)
    message.attach(body_part)

    # section 1 to attach file
    if type(attachments) is not list:
        attachments = [attachments]
        
    for path_to_file in attachments:
        with open(path_to_file,'rb') as file:
            # Attach the file with filename to the email
            message.attach(MIMEApplication(file.read(), Name='YourAIMasterpiece.png'))

    # section 2 for sending email
    print('Platform detected:', sys.platform)
    if sys.platform == 'win32':
        smtp_port = kwargs.get('smtp_port', 465)
        with smtplib.SMTP_SSL(smtp_server, smtp_port) as server:
            server.login(sender, "duup ajon shhz hlgk")
            server.sendmail(sender, recipient, message.as_string())
            server.quit()
    else:
        smtp_port = kwargs.get('smtp_port', 2525)
        with smtplib.SMTP(smtp_server, smtp_port) as server:
            server.starttls()
            server.login(sender, "duup ajon shhz hlgk")
            server.sendmail(sender, recipient, message.as_string())
            server.quit()
    print('Email sent!')

email_regex = re.compile(r"[^@]+@[^@]+\.[^@]+")

modalities = [
    "Computed Tomography (CT)", "Digital Breast Tomosynthesis", "Digital Mammography",
    "Magnetic Resonance Imaging (MRI)", "Micro-CT", "PET-CT", 
    "Positron Emission Tomography (PET)", "Radiography", "Ultrasonography"
]

organs = [
    "Aorta", "Bladder", "Bone", "Brain", "Breast", "Cervix", "Chest", "Colon",
    "Coronary Heart", "Ear", "Endometrium", "Esophagus", "Eye", "Head And Neck", "Kidney", "Liver", 
    "Lung", "Lymph Node", "Ovary", "Pancreas", "Pelvis", "Peripheral Arterial", "Phantom", 
    "Prostate", "Rectum", "Uterus"
]

artists = [
    "Arcimboldo", "Leonardo da Vinci", "Botticelli", "Vincent van Gogh", "René Magritte", "Pablo Picasso", "Claude Monet", "Rembrandt",
    "Michelangelo", "Raffaello", "Henri Matisse", "Giotto", "Paul Cezanne", "Gustav Klimt", "Giorgio De Chirico", "Katsushika Hokusai",
    "Jackson Pollock", "Edgar Degas", "Francisco Goya", "Edouard Manet", "Andy Warhol", "Tintoretto", "Wu Daozi", "Banksy",
    "Salvador Dalí", "Wassily Kandinsky", "Paul Gauguin", "Joan Miró", "Georges Seurat", "Paul Signac", "Pelizza da Volpedo",
    "Édouard Vuillard", "Marc Chagall", "Kazimir Malevich", "Piet Mondrian", "Jean-Michel Basquiat",
    "Frida Kahlo", "Artemisia Gentileschi", "Caravaggio", "El Greco", "Eugène Delacroix",
    "J.M.W. Turner"
]

st.write(' ')
st.markdown("<div style=\"text-align: center;\"> <img src=\"https://www.deeptracetech.com/images/graphical-verbose-logo-inline.png\" alt=\"DeepTrace Technologies logo\" width=\"208\" height=\"100\"></div>", 
            unsafe_allow_html=True)
st.html(
    body="""
        <style>
        h1 > div > a {display: none !important;};
        </style>
        <h1 style='text-align: center;'>Paint your medical image<br/>as a Great Master would!</h1>
    """
)
organ = st.selectbox('Organ', organs, index=None)
modality = st.selectbox('Modality', modalities, index=None)
style = st.selectbox('Style', sorted(artists), index=None)
receiver = st.text_input('Email Address (optional)', placeholder="example@mail.com", value='')

privacy = st.checkbox('I have read and understood the privacy statement.', help=privacy_statement)

image_descriptions = {
    "Computed Tomography (CT)": "CT scan",
    "Digital Breast Tomosynthesis": "3D breast tomosynthesis",
    "Digital Mammography": "digital mammography",
    "Magnetic Resonance Imaging (MRI)": "MRI scan",
    "Micro-CT": "micro-CT scan",
    "PET-CT": "PET-CT scan",
    "Positron Emission Tomography (PET)": "PET scan",
    "Radiography": "radiography",
    "Ultrasonography": "ultrasonography"
}

organ_descriptions = {
    "Anus": "of the anus, focusing on musculature and surrounding tissues",
    "Aorta": "of the aorta, capturing the main artery and branches",
    "Bladder": "of the bladder, emphasizing its structure",
    "Bone": "of the bone, showcasing its structure",
    "Brain": "of the brain, highlighting its structures",
    "Breast": "of the breast, focusing on its anatomy",
    "Cervix": "of the cervix, capturing its features",
    "Chest": "of the chest, showcasing thoracic cavity",
    "Colon": "of the colon, emphasizing its structure",
    "Coronary Heart": "of the coronary heart, highlighting coronary arteries",
    "Ear": "of the ear, focusing on outer, middle, and inner ear",
    "Endometrium": "of the endometrium, highlighting uterine lining",
    "Esophagus": "of the esophagus, emphasizing its structure",
    "Eye": "of the eye, capturing its features",
    "Head And Neck": "of the head and neck, showcasing anatomical structures",
    "Kidney": "of the kidney, focusing on internal structure",
    "Liver": "of the liver, highlighting its lobes",
    "Lung": "of the lung, showcasing respiratory anatomy",
    "Lymph Node": "of the lymph node, emphasizing anatomical features",
    "Ovary": "of the ovary, highlighting internal structures",
    "Pancreas": "of the pancreas, showcasing its structure",
    "Pelvis": "of the pelvis, focusing on bone structure",
    "Peripheral Arterial": "of the peripheral arterial system, highlighting vascular structures",
    "Phantom": "of a phantom, focusing on simulated features",
    "Prostate": "of the prostate, emphasizing its structure",
    "Rectum": "of the rectum, highlighting its anatomy",
    "Soft Tissues": "of the soft tissues, showcasing detailed anatomy",
    "Uterus": "of the uterus, emphasizing its structure"
}

if organ != None:
    organ_description = organ_descriptions[organ]
if modality != None:
    image_description = image_descriptions[modality]


artist_prompts = {
    'Leonardo da Vinci': f"A detailed {modality} of the {organ} in the intricate and detailed style of Leonardo da Vinci's anatomical drawings. The {modality} should accurately depict the structure of the {organ}, combining scientific precision with Renaissance artistry.",
    'Vincent van Gogh': f"A detailed {modality} of the {organ} in the style of Vincent van Gogh, characterized by swirling, vibrant colors and bold brushstrokes reminiscent of 'Starry Night'. The {modality} should accurately portray the {organ} structure while embracing van Gogh's artistic expression.",
    'Pablo Picasso': f"A detailed {modality} of the {organ} in the abstract and fragmented style of Pablo Picasso's Cubism. The {modality} should use geometric shapes and bold lines to depict the {organ} structure accurately, while reflecting Picasso's avant-garde artistic approach.",
    'Claude Monet': f"A detailed {modality} of the {organ} in the impressionistic style of Claude Monet. The {modality} should capture the {organ} structure with delicate brushstrokes and a soft color palette, maintaining accuracy in medical depiction and impressionistic artistry.",
    'Rembrandt': f"A detailed {modality} of the {organ} rendered in the dramatic and textured style of Rembrandt. The {modality} should emphasize the {organ} structure with deep shadows and highlights, while accurately representing medical details and Rembrandt's artistic flair.",
    'Michelangelo': f"A detailed {modality} of the {organ} sculpted in the grand and dynamic style of Michelangelo. The {modality} should highlight the {organ} structure with intricate details and a sense of movement, combining anatomical accuracy with Michelangelo's artistic vision.",
    'Raffaello': f"A detailed {modality} of the {organ} in the harmonious and balanced style of Raphael. The {modality} should depict the {organ} structure with clear lines and soft colors, maintaining precision in medical representation and Raphael's artistic harmony.",
    'Henri Matisse': f"A detailed {modality} of the {organ} in the vibrant and expressive style of Henri Matisse. The {modality} should use bold colors and fluid lines to depict the {organ} structure accurately, while capturing Matisse's expressive artistic approach.",
    'Paul Cezanne': f"A detailed {modality} of the {organ} in the structured and methodical style of Paul Cezanne. The {modality} should depict the {organ} structure with carefully arranged shapes and a harmonious color palette, reflecting Cezanne's methodical and anatomically accurate depiction.",
    'Gustav Klimt': f"A detailed {modality} of the {organ} in the ornate and decorative style of Gustav Klimt. The {modality} should use intricate patterns, gold leaf, and symbolic elements to depict the {organ} structure accurately, while embracing Klimt's decorative artistic vision.",
    'Jackson Pollock': f"A detailed {modality} of the {organ} in the chaotic and energetic style of Jackson Pollock's action painting. The {modality} should use splatters and drips of paint to create a dynamic image, while accurately representing the {organ} structure with medical precision.",
    'Edgar Degas': f"A detailed {modality} of the {organ} in the graceful and observational style of Edgar Degas. The {modality} should use delicate lines and soft pastels to depict the {organ} structure accurately, capturing Degas' anatomical sensitivity and artistic finesse.",
    'Francisco Goya': f"A detailed {modality} of the {organ} in the dark and dramatic style of Francisco Goya. The {modality} should use stark contrasts and a haunting atmosphere to depict the {organ} structure accurately, reflecting Goya's anatomically precise and emotionally intense artistic vision.",
    'Edouard Manet': f"A detailed {modality} of the {organ} in the bold and modern style of Edouard Manet. The {modality} should use strong contrasts and a focus on contemporary realism to depict the {organ} structure accurately, while embodying Manet's artistic realism and anatomical precision.",
    'Andy Warhol': f"A detailed {modality} of the {organ} in the pop art style of Andy Warhol. The {modality} should use bright colors, repetitive patterns, and a commercial aesthetic to depict the {organ} structure accurately, while capturing Warhol's distinctive artistic approach and anatomical details.",
    'Salvador Dalí': f"A detailed {modality} of the {organ} in the surreal and dreamlike style of Salvador Dalí. The {modality} should incorporate fantastical elements and imaginative twists, while accurately representing the {organ} structure with medical precision and Dalí's surreal artistic expression.",
    'Wassily Kandinsky': f"A detailed {modality} of the {organ} in the abstract and colorful style of Wassily Kandinsky. The {modality} should use vibrant hues and geometric shapes to depict the {organ} structure accurately, reflecting Kandinsky's abstract artistic vision and anatomical precision.",
    'Paul Gauguin': f"A detailed {modality} of the {organ} in the exotic and symbolic style of Paul Gauguin. The {modality} should use rich colors and simplified forms to depict the {organ} structure accurately, while capturing Gauguin's symbolic and anatomically stylized artistic approach.",
    'Joan Miró': f"A detailed {modality} of the {organ} in the playful and abstract style of Joan Miró. The {modality} should use whimsical shapes and bright colors to depict the {organ} structure accurately, while embodying Miró's abstract artistic expression and anatomical precision.",
    'Georges Seurat': f"A detailed {modality} of the {organ} in the pointillist style of Georges Seurat. The {modality} should use meticulously arranged dots of color to depict the {organ} structure accurately, while embracing Seurat's precise artistic technique and anatomical detail.",
    'Édouard Vuillard': f"A detailed {modality} of the {organ} in the intimate and decorative style of Édouard Vuillard. The {modality} should use soft colors and intricate patterns to depict the {organ} structure accurately, reflecting Vuillard's intimate artistic portrayal and anatomical precision.",
    'Marc Chagall': f"A detailed {modality} of the {organ} in the dreamlike and fantastical style of Marc Chagall. The {modality} should use floating figures and vivid colors to depict the {organ} structure accurately, capturing Chagall's dreamlike artistic expression and anatomical detail.",
    'Kazimir Malevich': f"A detailed {modality} of the {organ} in the minimalist and geometric style of Kazimir Malevich. The {modality} should use simple shapes and pure colors to depict the {organ} structure accurately, reflecting Malevich's minimalist artistic vision and anatomical precision.",
    'Piet Mondrian': f"A detailed {modality} of the {organ} in the neoplasticist style of Piet Mondrian. The {modality} should use a grid of black lines and primary colors to depict the {organ} structure accurately, embodying Mondrian's precise artistic composition and anatomical detail.",
    'Jean-Michel Basquiat': f"A detailed {modality} of the {organ} in the raw and expressive style of Jean-Michel Basquiat. The {modality} should use bold lines and graffiti-like elements to depict the {organ} structure accurately, capturing Basquiat's expressive artistic expression and anatomical detail.",
    'Frida Kahlo': f"A detailed {modality} of the {organ} in the deeply personal and symbolic style of Frida Kahlo. The {modality} should use vivid colors and emotional intensity to depict the {organ} structure accurately, reflecting Kahlo's emotionally charged artistic expression and anatomical detail.",
    'Artemisia Gentileschi': f"A detailed {modality} of the {organ} in the dramatic and powerful style of Artemisia Gentileschi. The {modality} should use strong contrasts and dynamic composition to depict the {organ} structure accurately, embodying Gentileschi's powerful artistic expression and anatomical detail.",
    'Caravaggio': f"A detailed {modality} of the {organ} in the intense and chiaroscuro style of Caravaggio. The {modality} should use stark lighting and dramatic shadows to depict the {organ} structure accurately, capturing Caravaggio's dramatic artistic vision and anatomical detail.",
    'El Greco': f"A detailed {modality} of the {organ} in the elongated and expressive style of El Greco. The {modality} should use distorted forms and vibrant colors to depict the {organ} structure accurately, reflecting El Greco's expressive artistic expression and anatomical detail.",
    'Eugène Delacroix': f"A detailed {modality} of the {organ} in the romantic and dynamic style of Eugène Delacroix. The {modality} should use bold colors and energetic brushstrokes to depict the {organ} structure accurately, embodying Delacroix's dramatic artistic expression and anatomical precision.",
    'J.M.W. Turner': f"A detailed {modality} of the {organ} in the atmospheric and luminous style of J.M.W. Turner. The {modality} should use soft, swirling colors and a sense of light and movement to depict the {organ} structure accurately, capturing Turner's evocative artistic vision and anatomical detail.",
    'Arcimboldo': f"A detailed {modality} of the {organ} in the imaginative and composite style of Arcimboldo. The {modality} should depict the {organ} structure using elements like fruits, vegetables, and other objects, while ensuring anatomical accuracy and Arcimboldo's whimsical artistic approach.",
    'Botticelli': f"A detailed {modality} of the {organ} in the elegant and flowing style of Botticelli. The {modality} should use delicate lines and graceful compositions to depict the {organ} structure accurately, capturing Botticelli's refined artistic expression and anatomical precision.",
    'René Magritte': f"A detailed {modality} of the {organ} in the surreal and thought-provoking style of René Magritte. The {modality} should incorporate unexpected elements and juxtapositions to depict the {organ} structure accurately, reflecting Magritte's surreal artistic vision and anatomical detail.",
    'Giotto': f"A detailed {modality} of the {organ} in the early Renaissance style of Giotto. The {modality} should use clear lines and solid forms to depict the {organ} structure accurately, reflecting Giotto's pioneering artistic approach and anatomical precision.",
    'Giorgio De Chirico': f"A detailed {modality} of the {organ} in the metaphysical style of Giorgio De Chirico. The {modality} should use eerie, dreamlike settings and classical elements to depict the {organ} structure accurately, capturing De Chirico's enigmatic artistic vision and anatomical detail.",
    'Katsushika Hokusai': f"A detailed {modality} of the {organ} in the ukiyo-e style of Katsushika Hokusai. The {modality} should use bold outlines and vibrant colors to depict the {organ} structure accurately, reflecting Hokusai's distinctive artistic style and anatomical precision.",
    'Tintoretto': f"A detailed {modality} of the {organ} in the dynamic and dramatic style of Tintoretto. The {modality} should use vigorous brushstrokes and intense lighting to depict the {organ} structure accurately, capturing Tintoretto's energetic artistic vision and anatomical detail.",
    'Wu Daozi': f"A detailed {modality} of the {organ} in the expressive and fluid style of Wu Daozi. The {modality} should use flowing lines and a sense of movement to depict the {organ} structure accurately, reflecting Wu Daozi's dynamic artistic approach and anatomical precision.",
    'Banksy': f"A detailed {modality} of the {organ} in the provocative and street art style of Banksy. The {modality} should use stenciled forms and bold statements to depict the {organ} structure accurately, capturing Banksy's distinctive artistic expression and anatomical detail.",
    'Paul Signac': f"A detailed {modality} of the {organ} in the pointillist style of Paul Signac. The {modality} should use meticulously arranged dots of color to depict the {organ} structure accurately, while embracing Signac's precise artistic technique and anatomical detail.",
    'Pelizza da Volpedo': f"A detailed {modality} of the {organ} in the Divisionist style of Pelizza da Volpedo. The {modality} should use small, distinct dots of color to depict the {organ} structure accurately, reflecting Pelizza da Volpedo's precise artistic technique and anatomical detail."
}


combined_prompts = {artist: prompt.replace("{image_description}", "{image_description}").replace("{organ_description}", "{organ_description}")
                    for artist, prompt in artist_prompts.items()}


prompt_lst = [organ, modality, style]

if (None not in prompt_lst) and ((receiver == '') or (privacy)):
    st.session_state.button_disabled = False
else:
    st.session_state.button_disabled = True

if st.session_state.button_clicked:
    st.session_state.button_disabled = True
    st.session_state.button_clicked = False
    col1, col2, col3, col4, col5 = st.columns(5)
    with col3:
        st.button('Generate', disabled=st.session_state.button_disabled)
    with st.spinner('Processing...'):
        img_name = modality+'_'+organ+'_'+style+'_'+datetime.now().strftime("%Y%m%d_%H%M%S")+'.png'

        print('Prompt list:', prompt_lst)
        #prompt = "Generate a painting of {a}, in the style of {b}. Emphasize all the anatomical details captured by a {c}".format(a=image_description, b=style, c=organ_description)
        #prompt = "Create a {modality} image {organ}, in the style of {style}. The image should maintain the recognizable features of a {modality}, such as the cross-sectional view and structures, while incorporating artistic textures, colors, and brushstrokes".format(modality=image_description, organ=organ_description, style=style)
        prompt = combined_prompts[style].format(image_description=image_description, organ_description=organ_description) + 'High Resolution.'
        print('Full prompt:', prompt)
        pipe = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.bfloat16, variant="fp16")
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)

        pipe.to("cuda")

        image = pipe(prompt=prompt, guidance_scale=0.0, num_inference_steps=10, height=1024, width=1024).images[0]
        print('Done!')
        image = add_annotation(image, f"{style}: \"{organ}\"", font_path="./Raleway-Italic.ttf", font_size=36, margin=[20, 10])
        image = add_annotation(image, f"{modality} (AI generated)", font_path="./Raleway-Medium.ttf", font_size=36, margin=[0, 10])
        image = add_annotation(image, f"powered by: DeepTrace Technologies srl, Italy ©", font_path="./Raleway-Medium.ttf", font_size=36, margin=[0, 10])
        image = add_annotation(image, f"info@deeptracetech.com", font_path="./Raleway-Medium.ttf", font_size=28, margin=[0, 10])
        image.save(img_name)
    
    st.image(image)
    
    # Upload image on google drive
    gauth = drive_authenticate()
    drive = GoogleDrive(gauth)
    drive_upload(drive, drive_folder_name, img_name)


    if receiver != '':
        # Update emails list
        drive_download(drive, drive_folder_name, drive_txt_name)

        # Read the file content and add a new line
        with open(drive_txt_name, 'a') as file:
            file.write('\n')
            file.write(f'{receiver} - {img_name}')
        drive_update(drive, drive_folder_name, drive_txt_name)

    #if receiver != '':
    #    send_email('Your AI Masterpiece is here!', mail_body, img_name, sender, receiver)
    
    try:
        os.remove(img_name)
    except:
        pass

    st.session_state.button_disabled = False

else:
    col1, col2, col3, col4, col5 = st.columns(5)
    with col3:
        st.button('Generate', on_click=on_button_click, disabled=st.session_state.button_disabled)