app.py

import spaces  # Must remain first
import os
import pickle
import numpy as np
import tensorflow as tf
import gradio as gr
from PIL import Image

# Prevent TF from hogging all memory if a GPU is present
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)

# Import existing architectures
from model_scratch import CNN_Encoder, RNN_Decoder
from model_pretrained import CNN_Encoder_Pretrained

# Hyperparameters matching your training configurations
EMBEDDING_DIM = 256
UNITS = 512

# Global model placeholders
models = {
    'scratch': {'encoder': None, 'decoder': None, 'tokenizer': None, 'input_size': 299},
    'pretrained': {'encoder': None, 'decoder': None, 'tokenizer': None, 'input_size': 224}
}

def load_models():
    """Initializes and builds architectures for both versions."""
    for model_type in ['scratch', 'pretrained']:
        checkpoint_dir = f"checkpoints-models/{model_type}_model"
        tokenizer_path = os.path.join(checkpoint_dir, 'tokenizer.pickle')

        if not os.path.exists(tokenizer_path):
            continue

        with open(tokenizer_path, 'rb') as f:
            tokenizer = pickle.load(f)
        vocab_size = len(tokenizer.word_index) + 1

        if model_type == 'scratch':
            encoder = CNN_Encoder(embedding_dim=EMBEDDING_DIM)
            input_size = 299
        else:
            encoder = CNN_Encoder_Pretrained(embedding_dim=EMBEDDING_DIM, backbone='resnet50v2', fine_tune=False)
            input_size = 224

        decoder = RNN_Decoder(embedding_dim=EMBEDDING_DIM, units=UNITS, vocab_size=vocab_size, rnn_type='lstm')

        # Dummy pass to build trackers
        dummy_img = tf.zeros((1, input_size, input_size, 3))
        features = encoder(dummy_img, training=False)
        hidden, carry = decoder.reset_state(batch_size=1)
        start_token = tokenizer.word_index.get('<start>', 1)
        decoder(tf.constant([[start_token]]), features, hidden, carry, training=False)

        # Load weights
        encoder.load_weights(os.path.join(checkpoint_dir, 'encoder_weights.weights.h5'))
        decoder.load_weights(os.path.join(checkpoint_dir, 'decoder_weights.weights.h5'))

        models[model_type] = {
            'encoder': encoder,
            'decoder': decoder,
            'tokenizer': tokenizer,
            'input_size': input_size
        }

def preprocess_image(img, model_type):
    input_size = models[model_type]['input_size']
    img = img.convert('RGB')
    img = img.resize((input_size, input_size))
    img_array = np.array(img, dtype=np.float32) / 255.0

    if model_type == 'pretrained':
        img_tensor = tf.keras.applications.resnet_v2.preprocess_input(img_array * 255.0)
    else:
        img_tensor = (img_array * 2.0) - 1.0

    return tf.expand_dims(img_tensor, 0)

def beam_search_decode(encoder, decoder, tokenizer, img_tensor, beam_width=3, max_len=60):
    features = encoder(img_tensor, training=False)
    hidden, carry = decoder.reset_state(batch_size=1)

    start_token = tokenizer.word_index.get('<start>', 1)
    end_token = tokenizer.word_index.get('<end>', 2)

    beams = [(0.0, [start_token], hidden, carry)]

    for _ in range(max_len):
        candidates = []
        all_done = True

        for score, seq, h, c in beams:
            if seq[-1] == end_token:
                candidates.append((score, seq, h, c))
                continue

            all_done = False
            dec_input = tf.expand_dims([seq[-1]], 0)
            predictions, next_h, next_c, _ = decoder(dec_input, features, h, c, training=False)

            log_probs = tf.nn.log_softmax(predictions[0]).numpy()
            top_indices = np.argsort(log_probs)[-beam_width:]

            for idx in top_indices:
                candidates.append((score + log_probs[idx], seq + [int(idx)], next_h, next_c))

        if all_done:
            break

        beams = sorted(candidates, key=lambda x: x[0], reverse=True)[:beam_width]

    best_seq = beams[0][1]
    output_words = [tokenizer.index_word.get(idx, '<unk>') for idx in best_seq 
                    if idx not in [start_token, end_token]]
    return " ".join(output_words)

# --- THE DYNAMIC FALLBACK ROUTING ---

# 1. The Core Inference (CPU Safe, No Decorator)
def core_inference(img, model_type, beam_width):
    if models['scratch']['encoder'] is None or models['pretrained']['encoder'] is None:
        load_models()

    if not models[model_type]['encoder']:
        return "Error: Model weights are missing."

    img_tensor = preprocess_image(img, model_type)
    cfg = models[model_type]

    return beam_search_decode(
        cfg['encoder'], cfg['decoder'], cfg['tokenizer'], 
        img_tensor, beam_width=int(beam_width)
    )

# 2. The GPU Wrapper
@spaces.GPU
def inference_gpu(img, model_type, beam_width):
    # This function claims the GPU. If quota is exceeded, spaces throws an error here.
    return core_inference(img, model_type, beam_width)

# 3. The Router Function
def inference_router(img, model_type, beam_width, force_cpu):
    # If the user explicitly checks the "Force CPU" box, bypass GPU entirely
    if force_cpu:
        gr.Info("Running on CPU by user request. This will take longer.")
        with tf.device('/CPU:0'):
            return core_inference(img, model_type, beam_width)
    
    # Try GPU first
    try:
        return inference_gpu(img, model_type, beam_width)
    except Exception as e:
        # If Hugging Face throws an error (e.g., Quota Exceeded), catch it and fallback
        gr.Warning("ZeroGPU quota exceeded or unavailable. Falling back to CPU processing. This will take longer.")
        with tf.device('/CPU:0'):
            return core_inference(img, model_type, beam_width)

# Gradio Web Canvas layout mapping
with gr.Blocks(title="Medical X-Ray Captioning System") as demo:
    gr.Markdown("# Medical X-Ray Captioning System (Scratch vs Pretrained)")
    gr.Markdown("Web interface running on Gradio engine for cloud hosting integration. **Now with Dynamic CPU Fallback.**")

    with gr.Row():
        with gr.Column():
            input_img = gr.Image(type="pil", label="Upload Chest X-Ray")
            model_select = gr.Radio(choices=["scratch", "pretrained"], value="scratch", label="Select Backbone Model")
            beam_slider = gr.Slider(minimum=1, maximum=7, step=1, value=3, label="Beam Search Width")
            
            # Added a manual override checkbox
            force_cpu_checkbox = gr.Checkbox(label="Force CPU Fallback (Check this if quota is empty to avoid wait queues)", value=False)
            
            btn = gr.Button("Generate Medical Diagnostics Report", variant="primary")

        with gr.Column():
            output_text = gr.Textbox(label="Generated Diagnostic Observations Summary", interactive=False, lines=8)

    btn.click(fn=inference_router, inputs=[input_img, model_select, beam_slider, force_cpu_checkbox], outputs=output_text)

if __name__ == "__main__":
    demo.launch()