added momdalities, organs and artists + buillt-in prompt engeneering

app.py

@@ -4,6 +4,7 @@ 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
 
@@ -11,9 +12,74 @@ if 'button_clicked' not in st.session_state:
 def on_button_click():
     st.session_state.button_clicked = True
 
-organ = st.selectbox('Organ', ['Brain', 'Thorax'], index=None)
-modality = st.selectbox('Modality', ['Magnetic Resonance Imaging', 'Computed Tomography'], index=None)
-style = st.selectbox('Style', ['Picasso', 'Van Gogh'], index=None)
+
+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)
+
+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."
+}
+
+
 
 prompt_lst = [organ, modality, style]
 
@@ -28,13 +94,12 @@ if st.session_state.button_clicked:
     st.session_state.button_clicked = False
     st.button('Submit', disabled=st.session_state.button_disabled)
     with st.spinner('Processing...'):
-        prompt = ','.join(prompt_lst) 
-        print(prompt)
+        print(prompt_lst)
 
+        prompt = artist_prompts[style]
         pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16")
         pipe.to("cuda")
 
-        prompt += " high resolution, photorealistic"
         image = pipe(prompt=prompt).images[0]
     
     st.image(image)