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"
]


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)": "a high-resolution cross-sectional image created by computed tomography (CT), capturing detailed structures of the",
    "DIGITAL BREAST TOMOSYNTHESIS": "a detailed 3D image generated by digital breast tomosynthesis, providing high-definition visualization of the",
    "DIGITAL MAMMOGRAPHY": "a clear, high-contrast image produced by digital mammography, highlighting the fine details and textures of the",
    "MAGNETIC RESONANCE IMAGING (MRI)": "a detailed, high-contrast magnetic resonance imaging (MRI) scan, providing in-depth visualization of the",
    "MICRO-CT": "a high-resolution micro-CT scan, capturing the minute structural details of the",
    "PET-CT": "a combined positron emission tomography-computed tomography (PET-CT) scan, offering both metabolic and anatomical information of the",
    "POSITRON EMISSION TOMOGRAPHY (PET)": "a detailed positron emission tomography (PET) scan, highlighting the metabolic and functional aspects of the",
    "RADIOGRAPHY": "a sharp, high-contrast radiography (X-ray) image, providing clear visualization of the",
    "ULTRASONOGRAPHY": "a detailed ultrasonography image, capturing the soft tissue structures of the"
}

# Descriptions for each organ
organ_descriptions = {
    "ANUS": "an anatomically precise depiction of the anus, focusing on detailed musculature and surrounding tissues",
    "AORTA": "a highly detailed image of the aorta, capturing the intricate structure of the main artery and its branches",
    "BLADDER": "a clear representation of the bladder, emphasizing the detailed structure of the organ and its surrounding tissues",
    "BONE": "an accurate image of the bone, showcasing its detailed structure, including cortical and trabecular patterns",
    "BRAIN": "a detailed visualization of the brain, highlighting its complex structures and intricate neural networks",
    "BREAST": "a precise image of the breast, focusing on its detailed anatomy and internal structures",
    "CERVIX": "a detailed depiction of the cervix, capturing its unique anatomical features and surrounding tissues",
    "CHEST": "a clear representation of the chest, showcasing the detailed anatomy of the thoracic cavity",
    "COLON": "an accurate image of the colon, emphasizing its detailed structure and surrounding tissues",
    "CORONARY HEART": "a precise visualization of the coronary heart, highlighting the intricate structure of the coronary arteries and heart muscles",
    "EAR": "a detailed image of the ear, focusing on the complex anatomy of the outer, middle, and inner ear",
    "ENDOMETRIUM": "an accurate depiction of the endometrium, highlighting the detailed layers of the uterine lining",
    "ESOPHAGUS": "a clear representation of the esophagus, emphasizing its detailed structure and surrounding tissues",
    "HEAD AND NECK": "a detailed visualization of the head and neck, showcasing the complex anatomical structures and tissues",
    "KIDNEY": "a precise image of the kidney, focusing on its detailed internal structure and surrounding tissues",
    "LIVER": "an accurate depiction of the liver, highlighting its detailed lobes and internal structures",
    "LUNG": "a clear representation of the lung, showcasing the detailed anatomy of the respiratory organ",
    "LYMPH NODE": "a detailed image of the lymph node, emphasizing its unique anatomical features and surrounding tissues",
    "OVARY": "an accurate depiction of the ovary, highlighting its detailed internal structures and surrounding tissues",
    "PANCREAS": "a clear image of the pancreas, showcasing its detailed structure and surrounding tissues",
    "PELVIS": "a precise representation of the pelvis, focusing on its detailed bone structure and surrounding tissues",
    "PERIPHERAL ARTERIAL": "a detailed depiction of the peripheral arterial system, highlighting its intricate vascular structures",
    "PHANTOM": "a detailed representation of a phantom, focusing on its simulated anatomical features for medical imaging purposes",
    "PROSTATE": "an accurate image of the prostate, emphasizing its detailed structure and surrounding tissues",
    "RECTUM": "a precise depiction of the rectum, highlighting its detailed anatomy and surrounding tissues",
    "SOFT TISSUES": "a clear representation of the soft tissues, showcasing the detailed anatomy and intricate structures",
    "UTERUS": "an accurate image of the uterus, emphasizing its detailed structure and surrounding tissues"
}

if organ != None:
    organ_description = organ_descriptions[organ.upper()]
if modality != None:
    image_description = image_descriptions[modality.upper()]


artist_prompts = {
    "Leonardo da Vinci": f"Generate a highly detailed and anatomically accurate {{image_description}} {{organ_description}} in the style of Leonardo da Vinci. Emphasize intricate anatomical details, precision, and a Renaissance aesthetic with detailed sketches and cross-hatching, akin to da Vinci's anatomical drawings.",
    "Vincent van Gogh": f"Create a {{image_description}} {{organ_description}} that is both medically accurate and stylistically reminiscent of Vincent van Gogh. Incorporate bold, expressive brushstrokes, vibrant colors, and swirling patterns that reflect van Gogh's emotional intensity and post-impressionist style.",
    "Pablo Picasso": f"Produce an accurate {{image_description}} {{organ_description}}, infused with the distinctive style of Pablo Picasso. Use cubist elements, fragmented forms, and abstract shapes to reinterpret the anatomical structure while maintaining medical accuracy.",
    "Claude Monet": f"Generate a {{image_description}} {{organ_description}} that combines medical precision with the impressionistic style of Claude Monet. Use soft, diffused colors and light, with a focus on blending and capturing the essence of the organ through light and shadow.",
    "Rembrandt": f"Create a {{image_description}} {{organ_description}} that reflects the Baroque style of Rembrandt. Emphasize dramatic lighting, deep contrasts, and a rich, textured appearance, with a focus on detailed anatomy and chiaroscuro effects.",
    "Michelangelo": f"Produce a {{image_description}} {{organ_description}} in the style of Michelangelo. Emphasize anatomical perfection, muscular detail, and a sculptural quality reminiscent of Michelangelo's work, with a Renaissance grandeur and precision.",
    "Raphael": f"Generate an anatomically accurate {{image_description}} {{organ_description}} with the refined elegance of Raphael. Use harmonious compositions, balanced proportions, and delicate, clear lines, reflecting the High Renaissance style.",
    "Henri Matisse": f"Create a {{image_description}} {{organ_description}} that blends medical accuracy with Henri Matisse's vibrant, Fauvist style. Use bold, flat colors, simplified forms, and a sense of dynamic movement to convey the anatomical structure.",
    "Paul Cezanne": f"Produce a {{image_description}} {{organ_description}} inspired by Paul Cezanne. Emphasize geometric simplification, structured forms, and a painterly approach with visible brushstrokes that convey depth and volume.",
    "Gustav Klimt": f"Generate a {{image_description}} {{organ_description}} in the opulent style of Gustav Klimt. Incorporate intricate patterns, gold leaf effects, and a decorative, almost mosaic-like quality, while ensuring anatomical accuracy.",
    "Jackson Pollock": f"Create a {{image_description}} {{organ_description}} with the dynamic, abstract expressionist style of Jackson Pollock. Use energetic splatters, drips, and layered colors to create an abstract yet medically precise representation.",
    "Edgar Degas": f"Produce an anatomically accurate {{image_description}} {{organ_description}} in the style of Edgar Degas. Use soft pastels, delicate lines, and a focus on capturing the subtle nuances and movement within the organ's structure.",
    "Francisco Goya": f"Generate a {{image_description}} {{organ_description}} reflecting the dramatic and dark style of Francisco Goya. Emphasize strong contrasts, intense emotion, and a somewhat somber tone, while maintaining medical accuracy.",
    "Edouard Manet": f"Create a {{image_description}} {{organ_description}} with the modern, impressionistic style of Edouard Manet. Use loose brushstrokes, a focus on light and shadow, and a slightly realist approach to anatomical detail.",
    "Andy Warhol": f"Produce an accurate {{image_description}} {{organ_description}} in the pop art style of Andy Warhol. Use bright, bold colors, repetitive patterns, and a graphic, almost commercial aesthetic to present the anatomical structure.",
    "Salvador Dalí": f"Generate a {{image_description}} {{organ_description}} with the surreal, dreamlike quality of Salvador Dalí. Incorporate unexpected elements, distorted forms, and a sense of fantastical realism, while ensuring anatomical precision.",
    "Wassily Kandinsky": f"Create a {{image_description}} {{organ_description}} in the abstract, expressionist style of Wassily Kandinsky. Use vibrant colors, geometric shapes, and a sense of rhythmic composition to depict the organ's structure.",
    "Paul Gauguin": f"Produce a {{image_description}} {{organ_description}} inspired by Paul Gauguin's post-impressionist style. Use bold, flat areas of color, strong outlines, and a somewhat primitive, symbolic approach to represent the anatomy.",
    "Joan Miró": f"Generate a {{image_description}} {{organ_description}} with the whimsical, abstract style of Joan Miró. Use playful shapes, bright colors, and a surreal, childlike quality, while maintaining accurate anatomical representation.",
    "Georges Seurat": f"Create an accurate {{image_description}} {{organ_description}} in the pointillist style of Georges Seurat. Use tiny dots of color to build up the image, focusing on light and color interactions to depict the anatomical details.",
    "Édouard Vuillard": f"Produce a {{image_description}} {{organ_description}} with the intimate, decorative style of Édouard Vuillard. Use soft colors, intricate patterns, and a focus on texture and detail, ensuring medical accuracy.",
    "Marc Chagall": f"Generate a {{image_description}} {{organ_description}} inspired by Marc Chagall's dreamlike, expressionist style. Use rich, vibrant colors, fantastical elements, and a poetic, almost narrative quality to depict the anatomy.",
    "Kazimir Malevich": f"Create an accurate {{image_description}} {{organ_description}} in the suprematist style of Kazimir Malevich. Use geometric shapes, abstract forms, and a minimalist approach, while ensuring clear anatomical representation.",
    "Piet Mondrian": f"Produce a {{image_description}} {{organ_description}} in the neoplasticism style of Piet Mondrian. Use clean lines, primary colors, and a grid-like composition to present the anatomical structure with precise accuracy.",
    "Jean-Michel Basquiat": f"Generate a {{image_description}} {{organ_description}} in the raw, expressive style of Jean-Michel Basquiat. Use bold lines, graffiti-like elements, and a chaotic yet powerful composition to depict the anatomy.",
    "Frida Kahlo": f"Create an anatomically accurate {{image_description}} {{organ_description}} with the emotive, surreal style of Frida Kahlo. Use rich colors, symbolic elements, and a deeply personal, almost autobiographical approach to represent the organ.",
    "Artemisia Gentileschi": f"Produce a {{image_description}} {{organ_description}} in the dramatic Baroque style of Artemisia Gentileschi. Emphasize strong contrasts, rich colors, and detailed anatomy, with a focus on powerful, emotional representation.",
    "Caravaggio": f"Generate a {{image_description}} {{organ_description}} with the intense, realistic style of Caravaggio. Use dramatic chiaroscuro, meticulous detail, and a lifelike quality to ensure anatomical precision.",
    "El Greco": f"Create a {{image_description}} {{organ_description}} inspired by El Greco's elongated forms and dramatic use of color. Emphasize fluid, almost ethereal shapes, and a strong sense of movement, while maintaining anatomical accuracy.",
    "Eugène Delacroix": f"Produce an accurate {{image_description}} {{organ_description}} with the romantic, dynamic style of Eugène Delacroix. Use vivid colors, energetic brushstrokes, and a sense of drama to depict the anatomical details.",
    "J.M.W. Turner": f"Generate a {{image_description}} {{organ_description}} in the atmospheric, impressionistic style of J.M.W. Turner. Use soft, diffused light, swirling colors, and a focus on capturing the essence and movement within the organ's structure."
}


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)