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('', 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('', 1) end_token = tokenizer.word_index.get('', 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, '') 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()