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) scan", "DIGITAL BREAST TOMOSYNTHESIS image", "DIGITAL MAMMOGRAPHY imge",
    "MAGNETIC RESONANCE IMAGING (MRI) scan", "MICRO-CT scan", "PET-CT scan", 
    "POSITRON EMISSION TOMOGRAPHY (PET) scan", "RADIOGRAPHY image", "ULTRASONOGRAPHY image"
]

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)

# artist_prompts = {
#     'Leonardo da Vinci': f"A detailed {organ} {modality} in the intricate and detailed style of Leonardo da Vinci's anatomical drawings. The {modality} should show the {organ} structure with medical precision, combining scientific accuracy with Renaissance artistry.",
#     'Vincent van Gogh': f"A detailed {organ} {modality} in the style of Van Gogh, with the swirling, vibrant colors and bold brushstrokes characteristic of his paintings. The {modality} should show the {organ} structure with medical precision, set against a background reminiscent of 'Starry Night.'",
#     'Pablo Picasso': f"A detailed {organ} {modality} depicted in the abstract and fragmented style of Pablo Picasso's Cubism. The {modality} should show the {organ} structure with medical precision, using geometric shapes and bold lines.",
#     'Claude Monet': f"A detailed {organ} {modality} in the impressionistic and soft-focus style of Claude Monet. The {modality} should show the {organ} structure with medical precision, using delicate brushstrokes and a subtle color palette.",
#     'Rembrandt': f"A detailed {organ} {modality} rendered in the dramatic and richly textured style of Rembrandt. The {modality} should show the {organ} structure with medical precision, using deep shadows and warm highlights.",
#     'Michelangelo': f"A detailed {organ} {modality} sculpted in the grand and dynamic style of Michelangelo. The {modality} should show the {organ} structure with medical precision, emphasizing intricate details and a sense of movement.",
#     'Raphael': f"A detailed {organ} {modality} in the harmonious and balanced style of Raphael. The {modality} should show the {organ} structure with medical precision, using clear lines and soft colors.",
#     'Henri Matisse': f"A detailed {organ} {modality} in the vibrant and expressive style of Henri Matisse. The {modality} should show the {organ} structure with medical precision, using bold colors and fluid lines.",
#     'Paul Cezanne': f"A detailed {organ} {modality} in the structured and methodical style of Paul Cezanne. The {modality} should show the {organ} structure with medical precision, using carefully arranged shapes and a harmonious color palette.",
#     'Gustav Klimt': f"A detailed {organ} {modality} in the ornate and decorative style of Gustav Klimt. The {modality} should show the {organ} structure with medical precision, using intricate patterns, gold leaf, and symbolic elements.",
#     'Jackson Pollock': f"A detailed {organ} {modality} in the chaotic and energetic style of Jackson Pollock's action painting. The {modality} should show the {organ} structure with medical precision, using splatters and drips of paint.",
#     'Edgar Degas': f"A detailed {organ} {modality} in the graceful and observational style of Edgar Degas. The {modality} should show the {organ} structure with medical precision, using delicate lines and soft pastels.",
#     'Francisco Goya': f"A detailed {organ} {modality} in the dark and dramatic style of Francisco Goya. The {modality} should show the {organ} structure with medical precision, using stark contrasts and a haunting atmosphere.",
#     'Edouard Manet': f"A detailed {organ} {modality} in the bold and modern style of Edouard Manet. The {modality} should show the {organ} structure with medical precision, using strong contrasts and a focus on contemporary realism.",
#     'Andy Warhol': f"A detailed {organ} {modality} in the pop art style of Andy Warhol. The {modality} should show the {organ} structure with medical precision, using bright colors, repetitive patterns, and a playful, commercial aesthetic.",
#     'Salvador Dalí': f"A detailed {organ} {modality} in the surreal and dreamlike style of Salvador Dalí. The {modality} should show the {organ} structure with medical precision, incorporating fantastical elements and imaginative twists.",
#     'Wassily Kandinsky': f"A detailed {organ} {modality} in the abstract and colorful style of Wassily Kandinsky. The {modality} should show the {organ} structure with medical precision, using vibrant hues and geometric shapes.",
#     'Paul Gauguin': f"A detailed {organ} {modality} in the exotic and symbolic style of Paul Gauguin. The {modality} should show the {organ} structure with medical precision, using rich colors and simplified forms.",
#     'Joan Miró': f"A detailed {organ} {modality} in the playful and abstract style of Joan Miró. The {modality} should show the {organ} structure with medical precision, using whimsical shapes and bright colors.",
#     'Georges Seurat': f"A detailed {organ} {modality} in the pointillist style of Georges Seurat. The {modality} should show the {organ} structure with medical precision, using tiny dots of color meticulously arranged.",
#     'Édouard Vuillard': f"A detailed {organ} {modality} in the intimate and decorative style of Édouard Vuillard. The {modality} should show the {organ} structure with medical precision, using soft colors and intricate patterns.",
#     'Marc Chagall': f"A detailed {organ} {modality} in the dreamlike and fantastical style of Marc Chagall. The {modality} should show the {organ} structure with medical precision, using floating figures and vivid colors.",
#     'Kazimir Malevich': f"A detailed {organ} {modality} in the minimalist and geometric style of Kazimir Malevich. The {modality} should show the {organ} structure with medical precision, using simple shapes and pure colors.",
#     'Piet Mondrian': f"A detailed {organ} {modality} in the neoplasticist style of Piet Mondrian. The {modality} should show the {organ} structure with medical precision, using a grid of black lines and primary colors.",
#     'Jean-Michel Basquiat': f"A detailed {organ} {modality} in the raw and expressive style of Jean-Michel Basquiat. The {modality} should show the {organ} structure with medical precision, using bold lines and graffiti-like elements.",
#     'Frida Kahlo': f"A detailed {organ} {modality} in the deeply personal and symbolic style of Frida Kahlo. The {modality} should show the {organ} structure with medical precision, using vivid colors and emotional intensity.",
#     'Artemisia Gentileschi': f"A detailed {organ} {modality} in the dramatic and powerful style of Artemisia Gentileschi. The {modality} should show the {organ} structure with medical precision, using strong contrasts and dynamic composition.",
#     'Caravaggio': f"A detailed {organ} {modality} in the intense and chiaroscuro style of Caravaggio. The {modality} should show the {organ} structure with medical precision, using stark lighting and dramatic shadows.",
#     'El Greco': f"A detailed {organ} {modality} in the elongated and expressive style of El Greco. The {modality} should show the {organ} structure with medical precision, using distorted forms and vibrant colors.",
#     'Eugène Delacroix': f"A detailed {organ} {modality} in the romantic and dynamic style of Eugène Delacroix. The {modality} should show the {organ} structure with medical precision, using bold colors and energetic brushstrokes.",
#     'J.M.W. Turner': f"A detailed {organ} {modality} in the atmospheric and luminous style of J.M.W. Turner. The {modality} should show the {organ} structure with medical precision, using soft, swirling colors and a sense of light and movement."
# }

artist_prompts = {
  "Leonardo da Vinci": f"""High-resolution {modality} of the {organ} with clear anatomy, rendered in Leonardo da Vinci's style with meticulous detail, soft shading, and realistic proportions on a parchment-like background (optional: focus on {organ} only, black background).""",
  "Vincent van Gogh": f"""High-resolution {modality} of the {organ} with clear anatomy, rendered in Van Gogh's style with thick swirling brushstrokes and vibrant colors (blues, yellows, greens) on a textured canvas (optional: focus on {organ} only, black background, slight glow).""",
  "Pablo Picasso": f"""High-resolution {modality} of the {organ} with a cubist style inspired by Pablo Picasso, featuring fragmented and geometric shapes with a limited color palette (optional: focus on {organ} only, black background).""",
  "Claude Monet": f"""High-resolution {modality} of the {organ} with a subtle, impressionistic style inspired by Claude Monet, featuring soft brushstrokes and blurred details with a focus on light and color (optional: focus on {organ} only, muted background).""",
  "Rembrandt": f"""High-resolution {modality} of the {organ} with a chiaroscuro style inspired by Rembrandt, featuring dramatic lighting, deep shadows, and rich details (optional: focus on {organ} only, dark background).""",
  "Michelangelo": f"""High-resolution {modality} of the {organ} with a detailed anatomical style inspired by Michelangelo, emphasizing musculature and form with a focus on realism (optional: focus on {organ} only, light background).""",
  "Raphael": f"""High-resolution {modality} of the {organ} with a harmonious and balanced style inspired by Raphael, featuring idealized proportions and soft lighting (optional: focus on {organ} only, muted background).""",
  "Henri Matisse": f"""High-resolution {modality} of the {organ} with a bold and colorful style inspired by Henri Matisse, featuring flat shapes and simplified forms (optional: focus on {organ} only, contrasting background).""",
  "Paul Cezanne": f"""High-resolution {modality} of the {organ} with a geometric and structured style inspired by Paul Cezanne, featuring geometric shapes and overlapping planes (optional: focus on {organ} only, muted background).""",
  "Gustav Klimt": f"""High-resolution {modality} of the {organ} with an ornamental and decorative style inspired by Gustav Klimt, featuring intricate patterns and gold accents (optional: focus on {organ} only, dark background).""",
  "Jackson Pollock": f"""High-resolution {modality} of the {organ} with an abstract expressionist style inspired by Jackson Pollock, featuring energetic gestures and dripped paint (optional: focus on {organ} only, contrasting background).""",
  "Edgar Degas": f"""High-resolution {modality} of the {organ} with a realistic and dynamic style inspired by Edgar Degas, focusing on movement and anatomy (optional: focus on {organ} in action, muted background).""",
  "Francisco Goya": f"""High-resolution {modality} of the {organ} with a dark and dramatic style inspired by Francisco Goya, featuring strong contrasts and expressive brushstrokes (optional: focus on {organ} only, dark background).""",
  "Edouard Manet": f"""High-resolution {modality} of the {organ} with a realistic and observational style inspired by Edouard Manet, focusing on light and shadow (optional: focus on {organ} only, muted background).""",
  "Andy Warhol": f"""High-resolution {modality} of the {organ} with a pop art style inspired by Andy Warhol, featuring bold colors and repeated patterns (optional: focus on {organ} only, contrasting background).""",
  "Salvador Dalí": f"""High-resolution {modality} of the {organ} with a surrealist style inspired by Salvador Dalí, featuring dreamlike imagery and distorted perspectives (optional: focus on {organ} only, dreamlike background).""",
  "Wassily Kandinsky": f"""High-resolution {modality} of the {organ} with an abstract style inspired by Wassily Kandinsky, featuring vibrant colors, geometric shapes, and dynamic compositions (optional: emphasize lines and circles, focus on {organ} only, black background).""",
  "Paul Gauguin": f"""High-resolution {modality} of the {organ} with a post-impressionist style inspired by Paul Gauguin, featuring bold colors, flat shapes, and exotic imagery (optional: focus on {organ} only, patterned background).""",
"Joan Miró": f"""High-resolution {modality} of the {organ} with a whimsical and playful style inspired by Joan Miró, featuring biomorphic shapes, vibrant colors, and childlike wonder (optional: focus on {organ} only, bright background).""",
"Georges Seurat": f"""High-resolution {modality} of the {organ} with a pointillist style inspired by Georges Seurat, featuring small dots of color that create a cohesive image when viewed from afar (optional: focus on {organ} only, muted background).""",
"Édouard Vuillard": f"""High-resolution {modality} of the {organ} with an intimate and decorative style inspired by Édouard Vuillard, featuring muted colors, flat patterns, and domestic scenes (optional: focus on {organ} only, patterned background).""",
"Marc Chagall": f"""High-resolution {modality} of the {organ} with a dreamlike and poetic style inspired by Marc Chagall, featuring floating figures, vibrant colors, and symbolic imagery (optional: focus on {organ} only, dreamlike background).""",
"Kazimir Malevich": f"""High-resolution {modality} of the {organ} with a suprematist style inspired by Kazimir Malevich, featuring geometric shapes (squares, circles, lines) in a limited color palette (optional: focus on {organ} only, black background).""",
"Piet Mondrian": f"""High-resolution {modality} of the {organ} with a neoplasticism style inspired by Piet Mondrian, featuring black lines, rectangular shapes, and primary colors (optional: focus on {organ} only, white background).""",
"Jean-Michel Basquiat": f"""High-resolution {modality} of the {organ} with a neo-expressionist style inspired by Jean-Michel Basquiat, featuring bold colors, raw energy, and symbolic imagery with text (optional: focus on {organ} only, contrasting background).""",
"Frida Kahlo": f"""High-resolution {modality} of the {organ} with a surrealist and self-portrait style inspired by Frida Kahlo, featuring bold colors, symbolic elements, and an exploration of identity (optional: focus on {organ} only, contrasting background).""",
"Artemisia Gentileschi": f"""High-resolution {modality} of the {organ} with a dramatic and chiaroscuro style inspired by Artemisia Gentileschi, featuring strong lighting, emotional expression, and realistic anatomy (optional: focus on {organ} only, dark background).""",
"Caravaggio": f"""High-resolution {modality} of the {organ} with a tenebrist style inspired by Caravaggio, featuring dramatic lighting, deep shadows, and intense realism (optional: focus on {organ} only, dark background).""",
"El Greco": f"""High-resolution {modality} of the {organ} with a mannerist style inspired by El Greco, featuring elongated figures, dramatic lighting, and emotional intensity (optional: focus on {organ} only, muted background).""",
"Eugène Delacroix": f"""High-resolution {modality} of the {organ} with a romantic style inspired by Eugène Delacroix, featuring movement, passionate expression, and dramatic lighting (optional: focus on {organ} in action, contrasting background).""",
"J.M.W. Turner": f"""High-resolution {modality} of the {organ} with a romantic and atmospheric style inspired by J.M.W. Turner, featuring light and shadow effects, dramatic landscapes, and a sense of movement (optional: focus on {organ} only, swirling background)."""
}



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 = artist_prompts[style]
        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)