app.py

import streamlit as st
import time
from diffusers import DiffusionPipeline
import matplotlib.pyplot as plt
import torch


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



modalities = [
    "COMPUTED TOMOGRAPHY (CT)", "DIGITAL BREAST TOMOSYNTHESIS", "DIGITAL MAMMOGRAPHY",
    "MAGNETIC RESONANCE IMAGING (MRI)", "MICRO-CT", "PET-CT", 
    "POSITRON EMISSION TOMOGRAPHY (PET)", "RADIOGRAPHY", "ULTRASONOGRAPHY"
]

modalities = [i.lower() for i in modalities]

organs = [
    "ANUS", "AORTA", "BLADDER", "BONE", "BRAIN", "BREAST", "CERVIX", "CHEST", "COLON",
    "CORONARY HEART", "EAR", "ENDOMETRIUM", "ESOPHAGUS", "HEAD AND NECK", "KIDNEY", "LIVER", 
    "LUNG", "LYMPH NODE", "OVARY", "PANCREAS", "PELVIS", "PERIPHERAL ARTERIAL", "PHANTOM", 
    "PROSTATE", "RECTUM", "SOFT TISSUES", "UTERUS"
]

organs = [i.lower() for i in organs]

artists = [
    "Leonardo da Vinci", "Vincent van Gogh", "Pablo Picasso", "Claude Monet", "Rembrandt",
    "Michelangelo", "Raphael", "Henri Matisse", "Paul Cezanne", "Gustav Klimt",
    "Jackson Pollock", "Edgar Degas", "Francisco Goya", "Edouard Manet", "Andy Warhol",
    "Salvador Dalí", "Wassily Kandinsky", "Paul Gauguin", "Joan Miró", "Georges Seurat",
    "É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.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.markdown("<h1 style='text-align: center;'>Paint your medical image as a great master would!</h1>", unsafe_allow_html=True)
#st.title("Paint your medical image as a great master would!")

organ = st.selectbox('Organ', organs, index=None)
modality = st.selectbox('Modality', modalities, index=None)
style = st.selectbox('Style', artists, index=None)

image_descriptions = {
    "COMPUTED TOMOGRAPHY (CT)": "CT scan",
    "DIGITAL BREAST TOMOSYNTHESIS": "3D breast tomosynthesis image",
    "DIGITAL MAMMOGRAPHY": "digital mammography image",
    "MAGNETIC RESONANCE IMAGING (MRI)": "MRI scan",
    "MICRO-CT": "micro-CT scan",
    "PET-CT": "PET-CT scan",
    "POSITRON EMISSION TOMOGRAPHY (PET)": "PET scan",
    "RADIOGRAPHY": "radiography image",
    "ULTRASONOGRAPHY": "ultrasonography image"
}

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",
    "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.upper()]
if modality != None:
    image_description = image_descriptions[modality.upper()]


artist_prompts = {
    "Leonardo da Vinci": "Generate a detailed {image_description} {organ_description} in the style of Leonardo da Vinci. Emphasize anatomical precision with Renaissance aesthetics.",
    "Vincent van Gogh": "Create a {image_description} {organ_description} in the style of Vincent van Gogh. Use bold brushstrokes and vibrant colors.",
    "Pablo Picasso": "Produce a {image_description} {organ_description} in the style of Pablo Picasso. Use cubist elements and abstract shapes.",
    "Claude Monet": "Generate a {image_description} {organ_description} in the style of Claude Monet. Use soft colors and diffused light.",
    "Rembrandt": "Create a {image_description} {organ_description} in the style of Rembrandt. Emphasize dramatic lighting and contrasts.",
    "Michelangelo": "Produce a {image_description} {organ_description} in the style of Michelangelo. Emphasize anatomical detail and sculptural quality.",
    "Raphael": "Generate an accurate {image_description} {organ_description} in the style of Raphael. Use harmonious compositions and delicate lines.",
    "Henri Matisse": "Create a {image_description} {organ_description} in the style of Henri Matisse. Use bold colors and simplified forms.",
    "Paul Cezanne": "Produce a {image_description} {organ_description} in the style of Paul Cezanne. Use geometric simplification and structured forms.",
    "Gustav Klimt": "Generate a {image_description} {organ_description} in the style of Gustav Klimt. Incorporate intricate patterns and a decorative quality.",
    "Jackson Pollock": "Create a {image_description} {organ_description} in the style of Jackson Pollock. Use energetic splatters and abstract layers.",
    "Edgar Degas": "Produce a {image_description} {organ_description} in the style of Edgar Degas. Use soft pastels and delicate lines.",
    "Francisco Goya": "Generate a {image_description} {organ_description} in the style of Francisco Goya. Emphasize strong contrasts and a somber tone.",
    "Edouard Manet": "Create a {image_description} {organ_description} in the style of Edouard Manet. Use loose brushstrokes and focus on light and shadow.",
    "Andy Warhol": "Produce a {image_description} {organ_description} in the style of Andy Warhol. Use bright colors and a graphic aesthetic.",
    "Salvador Dalí": "Generate a {image_description} {organ_description} in the style of Salvador Dalí. Incorporate surreal elements and distorted forms.",
    "Wassily Kandinsky": "Create a {image_description} {organ_description} in the style of Wassily Kandinsky. Use vibrant colors and geometric shapes.",
    "Paul Gauguin": "Produce a {image_description} {organ_description} in the style of Paul Gauguin. Use bold colors and strong outlines.",
    "Joan Miró": "Generate a {image_description} {organ_description} in the style of Joan Miró. Use playful shapes and bright colors.",
    "Georges Seurat": "Create an accurate {image_description} {organ_description} in the style of Georges Seurat. Use pointillism with tiny dots of color.",
    "Édouard Vuillard": "Produce a {image_description} {organ_description} in the style of Édouard Vuillard. Use soft colors and intricate patterns.",
    "Marc Chagall": "Generate a {image_description} {organ_description} in the style of Marc Chagall. Use rich colors and fantastical elements.",
    "Kazimir Malevich": "Create an accurate {image_description} {organ_description} in the style of Kazimir Malevich. Use geometric shapes and minimalism.",
    "Piet Mondrian": "Produce a {image_description} {organ_description} in the style of Piet Mondrian. Use clean lines and primary colors.",
    "Jean-Michel Basquiat": "Generate a {image_description} {organ_description} in the style of Jean-Michel Basquiat. Use bold lines and graffiti elements.",
    "Frida Kahlo": "Create an accurate {image_description} {organ_description} in the style of Frida Kahlo. Use rich colors and symbolic elements.",
    "Artemisia Gentileschi": "Produce a {image_description} {organ_description} in the style of Artemisia Gentileschi. Use strong contrasts and rich colors.",
    "Caravaggio": "Generate a {image_description} {organ_description} in the style of Caravaggio. Use dramatic lighting and meticulous detail.",
    "El Greco": "Create a {image_description} {organ_description} in the style of El Greco. Use elongated forms and dramatic color.",
    "Eugène Delacroix": "Produce a {image_description} {organ_description} in the style of Eugène Delacroix. Use vivid colors and energetic brushstrokes.",
    "J.M.W. Turner": "Generate a {image_description} {organ_description} in the style of J.M.W. Turner. Use soft light and swirling colors."
}




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:
    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
    st.button('Submit', disabled=st.session_state.button_disabled)
    with st.spinner('Processing...'):
        print(prompt_lst)

        prompt = combined_prompts[style].format(image_description=image_description, organ_description=organ_description)
        print(prompt)
        pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16")
        pipe.to("cuda")

        image = pipe(prompt=prompt).images[0]
    
    st.image(image)
    
    st.session_state.button_disabled = False
    

else:
    st.button('Submit', on_click=on_button_click, disabled=st.session_state.button_disabled)