Spaces:
Paused
Paused
| 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) |