app.py

# ------------------------------

# AI Multi‑Modal Assistant — Phase 2 (OCR Added)

# AI Multi‑Modal Assistant — Enhanced Version

# ------------------------------

import gradio as gr
from transformers import pipeline, pipeline as hf_pipeline, Wav2Vec2Processor, TFWav2Vec2ForCTC
from PIL import Image
import torch
from torchvision import models, transforms
import pandas as pd
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas
import io
from textwrap import wrap

import yake
import tempfile

import yake  # keyword extraction
import tempfile
from fpdf import FPDF
import time


import tensorflow as tf
import soundfile as sf
from langdetect import detect, DetectorFactory

DetectorFactory.seed = 0  # Consistent language detection

import os
import subprocess


# Ensure tesseract is available at runtime
try:
    subprocess.run(["tesseract", "--version"], check=True, stdout=subprocess.PIPE)
except (subprocess.CalledProcessError, FileNotFoundError):
    print("Installing tesseract runtime...")
    os.system("apt-get update -y && apt-get install -y tesseract-ocr")
        
import pytesseract  # <-- OCR


# ------------------------------
# 1. Load Models & Labels
# ------------------------------

# NLP pipelines
sentiment_model = pipeline(
    "sentiment-analysis",
    model="distilbert/distilbert-base-uncased-finetuned-sst-2-english",
)


# Summarization Model
summarizer_model = pipeline("summarization", model="facebook/bart-large-cnn")

# Image classification model
image_model = models.resnet50(pretrained=True)
image_model.eval()
preprocess = transforms.Compose(
    [
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ]
)

# Load ImageNet class labels

with open("imagenet_classes.txt", "r") as f:  # ensure this file is in your folder
    imagenet_labels = [s.strip() for s in f.readlines()]

# Keyword extraction
kw_extractor = yake.KeywordExtractor(lan="en", top=5)

# ------------------------------
# 2. Helper Functions
# ------------------------------

def analyze_text(text: str) -> dict:
    sentiment = sentiment_model(text)[0]
    summary = summarizer_model(
        text, max_length=min(len(text.split()) + 10, 50), min_length=5
    )[0]["summary_text"]
    keywords = [kw for kw, score in kw_extractor.extract_keywords(text)]
    return {
        "Sentiment": sentiment["label"],
        "Sentiment Score": round(sentiment["score"], 3),
        "Summary": summary,
        "Keywords": keywords,
    }

def analyze_image(image: Image.Image) -> dict:
    img_t = preprocess(image).unsqueeze(0)
    with torch.no_grad():
        outputs = image_model(img_t)
        class_idx = outputs.argmax().item()
    class_label = imagenet_labels[class_idx] if 0 <= class_idx < len(imagenet_labels) else f"Class index {class_idx}"
    return { "Predicted Class Index": class_idx, "Predicted Class Label": class_label }

def ocr_image(image: Image.Image) -> dict:
    """Extract text from uploaded image using Tesseract OCR."""
    text = pytesseract.image_to_string(image)
    return {"Extracted Text": text}

    if 0 <= class_idx < len(imagenet_labels):
        class_label = imagenet_labels[class_idx]
    else:
        class_label = f"Class index {class_idx}"
    return { "Predicted Class Index": class_idx, "Predicted Class Label": class_label }

def generate_pdf(results: dict) -> str:
    buffer = io.BytesIO()
    c = canvas.Canvas(buffer, pagesize=letter)
    width, height = letter
    c.setFont("Helvetica", 12)
    
    c.drawString(50, height - 50, "AI Multi-Modal Assistant Report")

    y = height -80      # Starting y position
     
    for key, value in results.items():
        # Convert any complex types to string
        if isinstance(value, (list, dict)):
            value = str(value)

        # Wrap long lines (max 90 chars per line)
        wrapped_lines = wrap(f"{key}: {value}", width=90)
        for line in wrapped_lines:
            c.drawString(50, y, line)
            y -= 18
            if y < 60:  # New page if needed
                c.showPage()
                c.setFont("Helvetica", 12)
                y = height - 50

        y -= 10  # Space between entries

    c.save()
    buffer.seek(0)

    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
        tmp.write(buffer.getvalue())
        tmp_path = tmp.name

    return tmp_path


    
# ------------------------------
# 3. Multi‑Modal Analysis Function
# ------------------------------
def analyze(input_data):
    # Handles both text and image correctly in Gradio
    if isinstance(input_data, str) and input_data.strip():
        return analyze_text(input_data)
    elif isinstance(input_data, dict) and "image" in input_data:
        return analyze_image(input_data["image"])
    elif isinstance(input_data, Image.Image):
        return analyze_image(input_data)
    else:
        return {"Error": "Please enter text or upload an image."}


# ------------------------------
# 4. Gradio UI Layout
# ------------------------------

with gr.Blocks() as demo:
    gr.Markdown("## AI Multi‑Modal Assistant")

    # ------------------ Image Analysis Tab ------------------

    with gr.Tab("Image Analysis"):
        image_input = gr.Image(type="pil", label="Upload an image for classification")
        analyze_image_button = gr.Button("Analyze Image")  # ✅ MOVED INSIDE the tab

        image_output = gr.JSON(label="Image Analysis Results")
        pdf_button_image = gr.Button("Download Report (PDF)")

        analyze_image_button.click(fn=analyze, inputs=image_input, outputs=image_output)
        pdf_button_image.click(
            fn=lambda x: generate_pdf(analyze(x)),
            inputs=image_input,
            outputs=gr.File(label="Download PDF Report"),
        )


    with gr.Tab("Text Analysis"):
        text_input = gr.Textbox(
            label="Enter text to analyze",
            placeholder="Type your text here...",
            lines=5
        )
        analyze_text_button = gr.Button("Analyze Text")
        text_output = gr.JSON(label="Text Analysis Results")
        pdf_button_text = gr.Button("Download Report (PDF)")

        # Text analysis events
        
        analyze_text_button.click(
            fn=analyze,
            inputs=text_input,
            outputs=text_output
        )
        pdf_button_text.click(
            fn=lambda x: generate_pdf(analyze(x)),
            inputs=text_input,
            outputs=gr.File(label="Download PDF Report")
        )

    
    # ------------------ OCR Tab ------------------
    with gr.Tab("OCR"):
        ocr_input = gr.Image(type="pil", label="Upload image for OCR")
        ocr_button = gr.Button("Extract Text (OCR)")
        ocr_output = gr.JSON(label="OCR Results")
        pdf_button_ocr = gr.Button("Download OCR PDF")

        ocr_button.click(fn=ocr_image, inputs=ocr_input, outputs=ocr_output)
        
        pdf_button_ocr.click(
            fn=lambda x: generate_pdf(x),
            inputs=ocr_output,
            outputs=gr.File(label="Download PDF Report"),
        )
        

    # --------------------------------
    # Enhanced Tab: Voice Input (Speech-to-Text + Translation + Language Detection)
    # ------------------------------
    
    # ------------------------------
    # PDF Generator with Unicode support
    # ------------------------------

# ------------------------------
# Auto-download DejaVuSans.ttf for all Unicode scripts
# ------------------------------
FONT_PATH = "DejaVuSans.ttf"
if not os.path.exists(FONT_PATH):
    import requests
    url = "https://github.com/dejavu-fonts/dejavu-fonts/raw/master/ttf/DejaVuSans.ttf"
    r = requests.get(url)
    with open(FONT_PATH, "wb") as f:
        f.write(r.content)

# ------------------------------
# PDF generator with full Unicode support
# ------------------------------
def generate_pdf(content_dict, output_file=None):
    pdf = FPDF()
    pdf.add_page()
    pdf.add_font("DejaVu", "", FONT_PATH, uni=True)
    pdf.set_font("DejaVu", size=12)

    for key, value in content_dict.items():
        pdf.set_font("DejaVu", "B", 12)
        pdf.cell(0, 10, f"{key}:", ln=True)
        pdf.set_font("DejaVu", "", 12)
        pdf.multi_cell(0, 10, value)
        pdf.ln(5)

    if not output_file:
        output_file = f"transcript_{int(time.time())}.pdf"
    pdf.output(output_file)
    return output_file

# ------------------------------
# TF-based STT (Wav2Vec2 for multilingual)
# ------------------------------
processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-large-xlsr-53")
stt_model = TFWav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-large-xlsr-53")

def transcribe_audio(file_path):
    speech, _ = sf.read(file_path)
    input_values = processor(speech, return_tensors="tf", sampling_rate=16000).input_values
    logits = stt_model(input_values).logits
    predicted_ids = tf.argmax(logits, axis=-1)
    transcription = processor.batch_decode(predicted_ids)[0]
    return transcription

# ------------------------------
# Language detection
# ------------------------------
def detect_language(text):
    try:
        lang_code = detect(text)
        # Convert language code to full language name (for PDF)
        import pycountry
        lang = pycountry.languages.get(alpha_2=lang_code)
        return lang.name if lang else lang_code
    except:
        return "unknown"

# ------------------------------
# TensorFlow translation (multilingual -> English)
# ------------------------------
model_name = "Helsinki-NLP/opus-mt-mul-en"
translator = pipeline("translation", model=model_name, framework="tf")

# ------------------------------
# Gradio App
# ------------------------------
with gr.Blocks() as demo:
    with gr.Tab("Voice Input"):
        gr.Markdown("### 🎤 Multilingual Speech-to-Text + Translation + PDF")
        audio_input = gr.Audio(sources=["microphone", "upload"], type="filepath", label="🎙️ Record or upload an audio file")
        transcribe_button = gr.Button("Transcribe, Detect & Translate")
        detected_lang_output = gr.Textbox(label="🌐 Detected Language", lines=1)
        transcription_output = gr.Textbox(label="🗣️ Original Transcription", lines=3)
        translation_output = gr.Textbox(label="🌍 Translated to English", lines=3)
        pdf_button_audio = gr.Button("Download Transcript (PDF)")

        def transcribe_detect_translate(audio_file):
            if not audio_file:
                return "", "Please record or upload an audio file.", ""

            # Step 1: Transcription
            original_text = transcribe_audio(audio_file)
            if not original_text.strip():
                return "", "No speech detected.", ""

            # Step 2: Language detection
            detected_lang = detect_language(original_text)

            # Step 3: Translate if not English
            translated_text = (
                translator(original_text)[0]["translation_text"]
                if detected_lang.lower() != "english"
                else original_text
            )

            return detected_lang, original_text, translated_text

        transcribe_button.click(
            fn=transcribe_detect_translate,
            inputs=audio_input,
            outputs=[detected_lang_output, transcription_output, translation_output]
        )

        pdf_button_audio.click(
            fn=lambda lang, orig, trans: generate_pdf({
                "Detected Language": lang,
                "Original Transcription": orig,
                "Translated to English": trans
            }),
            inputs=[detected_lang_output, transcription_output, translation_output],
            outputs=gr.File(label="Download PDF Report")
        )

demo.launch()