app.py

import streamlit as st
import tensorflow as tf
from tensorflow.keras.layers import InputLayer
from PIL import Image
import numpy as np
# from .modules.tools import Analysing_image, text_to_speech
from gtts import gTTS
from transformers import BioGptTokenizer, AutoModelForCausalLM,  pipeline, TFAutoModel
from deep_translator import GoogleTranslator, MyMemoryTranslator
import tempfile
import os
# Set the page configuration

st.set_page_config(page_title="Medical Image Classifier & Chatbot", layout="wide")
st.title("Medical Image Classifier & Chatbot")
st.sidebar.header("Medical Analysis")
keys={'0':'Cyst', '1':'Normal', '2':'Stone', '3':'Tumor'}

def translate_text(text, source='auto', target='es'):
    translators = [GoogleTranslator, MyMemoryTranslator]
    
    for translator in translators:
        try:
            translation = translator(source=source, target=target).translate(text)
            print(f"{translator.__name__}: {translation}")
            return translation
        except Exception as e:
            print(f"{translator.__name__} failed: {e}")
    
    print("All translators failed. No translation found.")
    return None

def Analysing_image(st, model, image_file):
    try:
        # Open and display the image
        image = Image.open(image_file)
        st.image(image, caption="Uploaded Image", use_container_width=True)

        # Preprocess the image:
        # Ensure it is resized to the input dimensions your model expects (150x150 in this example)
        img_resized = image.resize((150, 150))
        img_array = np.array(img_resized).astype('float32') / 255.0

        # If the image is not 3 channels (some images might be grayscale), repeat channels if necessary
        if img_array.ndim == 2:
            img_array = np.stack((img_array,)*3, axis=-1)
        elif img_array.shape[2] == 1:
            img_array = np.concatenate([img_array]*3, axis=-1)

        # Add batch dimension
        img_batch = np.expand_dims(img_array, axis=0)

        # Run prediction
        predictions = model.predict(img_batch)

        st.write("Prediction probabilities:", predictions)

        # Assuming a multi-class classification where the class with the highest probability is selected:
        predicted_class = np.argmax(predictions, axis=1)
        st.write("Prediction :", keys[str(predicted_class[0])])
        return keys[str(predicted_class[0])]
    except:
        return None
    
# Function to convert text to speech
def text_to_speech(text):
    tts = gTTS(text)
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
    tts.save(temp_file.name)
    return temp_file.name

@st.cache_resource  # Cache the model to avoid reloading on every interaction
def load_generator():
    # Use a medical-specific model like BioGPT or a general-purpose model like GPT-2
    tokenizer = BioGptTokenizer.from_pretrained("microsoft/BioGPT")
    model = AutoModelForCausalLM.from_pretrained("microsoft/BioGPT")
    return model, tokenizer

@st.cache_resource  # Cache the model to avoid reloading on every interaction
def load_summarizer():
    # Use a summarization model like "facebook/bart-large-cnn"
    summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
    return summarizer

# generator, tokenizer = load_generator()
generator = pipeline("text2text-generation", model="EleutherAI/gpt-neo-125M")
summarizer = load_summarizer()
translator = pipeline("translation_en_to_de", model="Helsinki-NLP/opus-mt-en-de")

image_file = st.sidebar.file_uploader("Upload an Image (.jpg, .jpeg, .png)", type=["jpg", "jpeg", "png"])
# Create tabs for each functionality
tab1, tab2, tab3, tab4 = st.tabs(["Classification", "Chatbot", "Translation & Summary", "Audio"])
predict_class = None
Summary = None

with tab1:
    # st.write("Upload your finalized model and an image to classify.")
    # Sidebar for uploading files   
    model_file = st.sidebar.file_uploader("Upload your Keras model (.h5 file)", type=["h5"])
   # Check if a model has been uploaded
    if model_file is not None:
        # Save the uploaded model file to disk
        with open("uploaded_model.h5", "wb") as f:
            f.write(model_file.getbuffer())
        st.sidebar.success("Model uploaded successfully!")
        
        # Attempt to load the model
        try:
            model = tf.keras.models.load_model("uploaded_model.h5")
            st.sidebar.info("Model loaded successfully!")
        except Exception as e:
            st.sidebar.error("Error loading model: " + str(e))
            st.stop()
        
        # Check if an image has been uploaded
        if image_file is not None:
           predict_class = Analysing_image(st, model, image_file)
        else:
            st.info("Please upload an image to classify.")
    else:
        st.info("Using Pretrained model")
        model = tf.keras.models.load_model("./models/medical_classifier/medical_classifier.h5")
        # model = TFAutoModel.from_pretrained('./models/medical_classifier')
        # config = model.config
        # if 'batch_shape' in config:
        #     config['input_shape'] = config.pop('batch_shape')[1:]  # Remove batch size
            
        # # Create the input layer with the corrected configuration
        # input_layer = InputLayer(**model.config)
        # # Rebuild the model (if necessary)
        # model.build(input_shape=model.config['input_shape'])

        if image_file is not None:
            predict_class = Analysing_image(st, model, image_file)
        else:
            st.info("Please upload an image to classify.")
# -----------------------------------------------------------------------------------------------

with tab2:
    if predict_class is not None:
        # # Create a prompt for the model
        if predict_class == 'Normal':
            prompt = f"What does it mean when the doctor say my IRM is Normal ?"
        else:
            if predict_class == "Cyst":
                prompt = f"What is brain cystic lesion ?"
            else:
                prompt = f"What is brain {predict_class} ?"
        # # Generate text using the Hugging Face model
        # with st.spinner("Generating description..."):
        #     inputs = tokenizer(prompt, return_tensors="pt")
        #     # Générer du texte
        #     output = generator.generate(
        #         inputs["input_ids"],
        #         max_length=300,  # Longueur maximale du texte généré
        #         num_return_sequences=1,  # Nombre de séquences à générer
        #         no_repeat_ngram_size=2,  # Éviter la répétition de phrases
        #         top_k=50,  # Contrôle la diversité du texte
        #         top_p=0.95,  # Contrôle la qualité du texte
        #         temperature=0.7,  # Contrôle la créativité du texte
        #     )
        
        #     # Décoder et afficher le texte généré
        #     # generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
        #     outputs = tokenizer.decode(output[0], skip_special_tokens=True)
        #     # output = generator(prompt, max_length=200, num_return_sequences=1)
        #     # st.session_state.chat_response = output[0]['generated_text']
        #     # st.session_state.chat_response = outputs

        #     # Decode the generated text, *excluding* the input prompt
        #     generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
        #     # Find the length of the original prompt
        #     prompt_length = len(tokenizer.decode(inputs["input_ids"][0], skip_special_tokens=True))
        #     # Extract the generated part of the text
        #     generated_response = generated_text[prompt_length:]
        #     st.session_state.chat_response = generated_response
        
        # medical_term = "myocardial infarction"
        # medical_term = prompt
        # description = generator(f"Explain the medical term: {medical_term}")

        user_input = st.text_area("Enter your prompt:", prompt)
        # Slider for controlling the length of the generated text
        max_length = st.slider("Max length of generated text", 50, 500, 100)   
        # Button to generate text
        if st.button("Generate Text"):
            with st.spinner("Generating text..."):
                # Generate text using the model
                output = generator(user_input, max_length=max_length, num_return_sequences=1)
                generated_text = output[0]['generated_text']
                # Display the generated text
                st.subheader("Description:")
                st.write(generated_text)
                # st.session_state.chat_response = outputs
                st.session_state.chat_response = generated_text
        
        # st.session_state.chat_response = outputs
        # st.session_state.chat_response = generated_text
        # Display the generated description
        # st.subheader("Generated Description:")
        # st.write(prompt, '--')
        # st.write(st.session_state.chat_response)


# Add translation and summary functionality
with tab3:
    st.header("Translation to German & Summary")
    if 'chat_response' in st.session_state and st.session_state.chat_response:
        medical_terms = st.session_state.chat_response
        # Translate to German
        # translator = GoogleTranslator(source='en', target='de')
        # german_translation = translator.translate(medical_terms)
        # Translate the text
        # translation = translate_text(medical_terms, target='de')
        translation = translator(medical_terms)
        outputs = 'Unable to translate, please Retry ...'
        if translation:
            outputs = translation[0]['translation_text']
            # st.subheader("Final Translation:")
            # st.write(translation)
        else:
            st.info("Unable to translate the text. Please try to refresh")
            if st.button('Refresh'):
                outputs = translator(prompt, target='de')
                outputs = outputs[0]['translation_text']

        # Create summary (simple example - in practice you might want to use a more sophisticated summarization method)
        # summary = " ".join(st.session_state.chat_response.split()[:30]) + "..."
        st.write("German Translation:", outputs)

        if st.button("Generate Summary"):
            if outputs != 'Unable to translate, please Retry ...':
                # Create a prompt for the model
                # prompt = f"{predict_class} is a medical condition that refers to "
                prompt = outputs
                # Generate a summary using the Hugging Face model
                with st.spinner("Generating summary..."):
                    # Summarize the prompt (you can adjust max_length and min_length)
                    summary = summarizer(prompt, max_length=85, min_length=60, do_sample=False)
                # Display the generated summary
                st.subheader("Generated Summary:")
                st.write(summary[0]['summary_text'])
            else:
                st.warning("Please enter a medical term.")
       
    else:
        st.info("No chatbot response available for translation and summary.")


# Add audio functionality
with tab4:
    st.header("Audio Output")
    if 'chat_response' in st.session_state and st.session_state.chat_response:
        # Convert chatbot response to audio
        audio_file = text_to_speech(st.session_state.chat_response)
        st.audio(audio_file)
        # Clean up temp file
        with open(audio_file, "rb") as file:
            btn = st.download_button(
                label="Download audio",
                data=file,
                file_name="chat_response.mp3",
                mime="audio/mpeg"
            )
        os.unlink(audio_file)
    else:
        st.info("No chatbot response available for audio conversion.")



print("Streamlit app updated with translation, summarization, and audio features.")