import gradio as gr
from PIL import Image
from transformers import pipeline
import scipy.io.wavfile as wavfile
import numpy as np

caption_image = pipeline("image-to-text", model="Salesforce/blip-image-captioning-large")

Narrator = pipeline("text-to-speech", model="kakao-enterprise/vits-ljs")


def generate_audio(text):
    # Generate speech from the input text using the Narrator (VITS model)
    Narrated_Text = Narrator(text)
    
    # Extract the audio data and sampling rate
    audio_data = np.array(Narrated_Text["audio"][0])
    sampling_rate = Narrated_Text["sampling_rate"]
    
    # Save the generated speech as a WAV file
    wavfile.write("generated_audio.wav", rate=sampling_rate, data=audio_data)
    
    # Return the filename of the saved audio file
    return "generated_audio.wav"


def caption_my_image(pil_image):
    # Use BLIP to generate a text description of the input image
    semantics = caption_image(images=pil_image)[0]["generated_text"]
    
    # Generate audio from the text description
    return generate_audio(semantics)


# Main functionality tab
main_tab = gr.Interface(
    fn=caption_my_image,
    inputs=[gr.Image(label="Select Image", type="pil")],
    outputs=[gr.Audio(label="Generated Audio")],
    title="Image Audio Captioning App",
    description="This application provides audio descriptions for images.."
)

# Information tab
info_tab = gr.Markdown("""
# Image Audio Captioning App

### Purpose
This application is designed to assist visually impaired users by providing audio descriptions of images. It can also be used in various scenarios such as creating audio captions for educational materials, enhancing accessibility for digital content, and more.

### How to Use
- **Step 1:** Click on the 'Select Image' button to upload an image.
- **Step 2:** Wait for the application to generate the audio description.
- **Step 3:** Listen to the generated audio file.

### Limits
- The quality of the description depends on the image clarity and content.
- The application might not work well with images that have complex scenes or unclear subjects.
- Audio generation time may vary depending on the input image size and content.

### Note
- Ensure the uploaded image is clear and well-defined for the best results.
- This app is a prototype and may have limitations in real-world applications.
""")

# Combine both tabs into a single app
demo = gr.TabbedInterface(
    [main_tab, info_tab],
    tab_names=["Main", "Information"]
)

demo.launch()