import streamlit as st
from streamlit_extras.switch_page_button import switch_page


translations = {
'en': {'title': 'SegGPT',
    'original_tweet': 
       """
       [Original tweet](https://x.com/mervenoyann/status/1773056450790666568) (March 27, 2024)
       """,
    'tweet_1':
        """
        SegGPT is a vision generalist on image segmentation, quite like GPT for computer vision ✨  
        It comes with the last release of 🤗 Transformers 🎁    
        Technical details, demo and how-to's under this! 
        """,
    'tweet_2':
        """
        SegGPT is an extension of the <a href='Painter' target='_self'>Painter</a> where you speak to images with images: the model takes in an image prompt, transformed version of the image prompt, the actual image you want to see the same transform, and expected to output the transformed image.  
        <br>
        SegGPT consists of a vanilla ViT with a decoder on top (linear, conv, linear). The model is trained on diverse segmentation examples, where they provide example image-mask pairs, the actual input to be segmented, and the decoder head learns to reconstruct the mask output. 👇🏻 
        """,
    'tweet_3':
        """
        This generalizes pretty well!  
        The authors do not claim state-of-the-art results as the model is mainly used zero-shot and few-shot inference. They also do prompt tuning, where they freeze the parameters of the model and only optimize the image tensor (the input context). 
        """,
    'tweet_4':
        """
        Thanks to 🤗 Transformers you can use this model easily! See [here](https://t.co/U5pVpBhkfK).
        """,
    'tweet_5':
        """
        I have built an app for you to try it out. I combined SegGPT with Depth Anything Model, so you don't have to upload image mask prompts in your prompt pair 🤗  
        Try it [here](https://t.co/uJIwqJeYUy). Also check out the [collection](https://t.co/HvfjWkAEzP).
        """,
    'ressources':
        """
        Ressources:   
        [SegGPT: Segmenting Everything In Context](https://arxiv.org/abs/2304.03284) 
        by Xinlong Wang, Xiaosong Zhang, Yue Cao, Wen Wang, Chunhua Shen, Tiejun Huang (2023)  
        [GitHub](https://github.com/baaivision/Painter)
        """
      },
'fr': {
    'title': 'SegGPT',
    'original_tweet': 
       """
       [Tweet de base](https://x.com/mervenoyann/status/1773056450790666568) (en anglais) (27 mars 2024)
       """,
    'tweet_1':
        """
        SegGPT est un modèle généraliste de vision pour la segmentation d'images;  c'est un peu comme le GPT pour la vision par ordinateur ✨.  
        Il est intégré à la dernière version de 🤗 Transformers 🎁    
        Détails techniques, démonstrations et manières de l'utiliser ci-dessous ! 
        """,
    'tweet_2':
        """
        SegGPT est une extension de <a href='Painter' target='_self'>Painter</a> où vous parlez aux images avec des images : le modèle reçoit une image, une version transformée de l'image, l'image réelle que vous voulez voir avec la même transformation, et est censé produire l'image transformée.  
        <br>
        SegGPT consiste en un ViT standard surmonté d'un décodeur (couche linéaire, convolution, couche linéaire). Le modèle est entraîné sur divers exemples de segmentation, où les auteurs fournissent des paires image-masque, l'entrée réelle à segmenter, et la tête du décodeur apprend à reconstruire la sortie du masque. 👇🏻         """,
    'tweet_3':
        """
        Cela se généralise assez bien !  
        Les auteurs ne prétendent pas obtenir des résultats de pointe, car le modèle est principalement utilisé pour l'inférence zéro-shot et few-shot. Ils effectuent également un prompt tuning, où ils gèlent les paramètres du modèle et optimisent uniquement le tenseur d'image (le contexte d'entrée). 
        """,
    'tweet_4':
        """ 
        Grâce à 🤗 Transformers, vous pouvez utiliser ce modèle facilement ! Voir [ici] (https://t.co/U5pVpBhkfK).
        """,
    'tweet_5':
        """
        J'ai créé une application pour que vous puissiez l'essayer. J'ai combiné SegGPT avec Depth Anything Model, de sorte que vous n'avez pas besoin de télécharger des masques d'images dans votre paire de prompt 🤗.  
        Essayez-le [ici](https://t.co/uJIwqJeYUy). Consultez également la [collection](https://t.co/HvfjWkAEzP).
        """,
    'ressources':
        """
        Ressources :   
        [SegGPT: Segmenting Everything In Context](https://arxiv.org/abs/2304.03284) 
        de Xinlong Wang, Xiaosong Zhang, Yue Cao, Wen Wang, Chunhua Shen, Tiejun Huang (2023)  
        [GitHub](https://github.com/baaivision/Painter)
        """
    }
}    


def language_selector():
    languages = {'EN': '🇬🇧', 'FR': '🇫🇷'}
    selected_lang = st.selectbox('', options=list(languages.keys()), format_func=lambda x: languages[x], key='lang_selector')
    return 'en' if selected_lang == 'EN' else 'fr'

left_column, right_column = st.columns([5, 1])

# Add a selector to the right column
with right_column:
    lang = language_selector()

# Add a title to the left column
with left_column:
    st.title(translations[lang]["title"])
    
st.success(translations[lang]["original_tweet"], icon="ℹ️")
st.markdown(""" """)

st.markdown(translations[lang]["tweet_1"], unsafe_allow_html=True)
st.markdown(""" """)

st.image("pages/SegGPT/image_1.jpeg", use_container_width=True)
st.markdown(""" """)

st.markdown(translations[lang]["tweet_2"], unsafe_allow_html=True)
st.markdown(""" """)

st.image("pages/SegGPT/image_2.jpg", use_container_width=True)
st.markdown(""" """)

st.markdown(translations[lang]["tweet_3"], unsafe_allow_html=True)
st.markdown(""" """)

st.image("pages/SegGPT/image_3.jpg", use_container_width=True)
st.markdown(""" """)

st.markdown(translations[lang]["tweet_4"], unsafe_allow_html=True)
st.markdown(""" """)

st.image("pages/SegGPT/image_4.jpeg", use_container_width=True)
st.markdown(""" """)

with st.expander ("Code"):
    st.code("""
    import torch 
    from transformers import SegGptImageProcessor, SegGptForImageSegmentation 
    
    image_processor = SegGptImageProcessor.from_pretrained("BAAI/seggpt-vit-large") 
    model = SegGptForImageSegmentation.from_pretrained("BAAI/seggpt-vit-large") 
    
    inputs = image_processor( 
        images=image_input, 
        prompt_images=image_prompt, 
        prompt_masks=mask_prompt, 
        num_labels=10, 
        return_tensors="pt") 
    
    with torch.no_grad(): 
        outputs = model(**inputs) 
    
    target_sizes = [image_input.size[::-1]] 
    mask = image_processor.post_process_semantic_segmentation(outputs, target_sizes, num_labels=10)[0]
    """)
st.markdown(""" """)

st.markdown(translations[lang]["tweet_5"], unsafe_allow_html=True)
st.markdown(""" """)

st.image("pages/SegGPT/image_5.jpeg", use_container_width=True)
st.markdown(""" """)

st.info(translations[lang]["ressources"], icon="📚")  

st.markdown(""" """)
st.markdown(""" """)
st.markdown(""" """)
col1, col2, col3= st.columns(3)
with col1:
    if lang == "en":
        if st.button('Previous paper', use_container_width=True):
            switch_page("Painter")
    else:
        if st.button('Papier précédent', use_container_width=True):
            switch_page("Painter")
with col2:
    if lang == "en":
        if st.button("Home", use_container_width=True):
            switch_page("Home")
    else:
        if st.button("Accueil", use_container_width=True):
            switch_page("Home")
with col3:
    if lang == "en":
        if st.button("Next paper", use_container_width=True):
            switch_page("Grounding DINO")
    else:
        if st.button("Papier suivant", use_container_width=True):
            switch_page("Grounding DINO")