Spaces:
No application file
No application file
| import os | |
| import sys | |
| import gradio as gr | |
| import torch | |
| from PIL import Image | |
| import nltk | |
| # Download NLTK data required for metrics/vocabulary on Hugging Face Spaces | |
| try: | |
| nltk.data.find('tokenizers/punkt') | |
| except LookupError: | |
| nltk.download('punkt', quiet=True) | |
| nltk.download('wordnet', quiet=True) | |
| nltk.download('omw-1.4', quiet=True) | |
| nltk.download('punkt_tab', quiet=True) | |
| # Add project root to path to allow importing from src | |
| sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) | |
| from src.inference.pipeline import CaptionPipeline | |
| # Global pipeline instance | |
| pipeline = None | |
| def load_model(): | |
| global pipeline | |
| if pipeline is None: | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| try: | |
| pipeline = CaptionPipeline.from_checkpoint('checkpoints', device=device) | |
| return f"✅ Model loaded successfully on **{device}**." | |
| except Exception as e: | |
| return f"❌ Failed to load model: {e}\n\nMake sure you have trained the model and `checkpoints/` exists." | |
| return f"✅ Model ready." | |
| def generate_caption(image, strategy, beam_size, max_length, temperature): | |
| if image is None: | |
| return "⚠️ Please upload an image first." | |
| if pipeline is None: | |
| return "❌ Model not loaded properly." | |
| try: | |
| # Generate caption using the pipeline | |
| result = pipeline.generate( | |
| image, | |
| strategy="beam" if strategy == "Beam Search" else "greedy", | |
| beam_size=int(beam_size), | |
| max_length=int(max_length), | |
| temperature=float(temperature) | |
| ) | |
| print(f"Generated caption: {result['caption']}") | |
| return f"Caption: {result['caption']}\nConfidence Score: {result['score']:.4f}" | |
| except Exception as e: | |
| print(f"Error: {e}") | |
| return f"Error generating caption: {e}" | |
| # Build Gradio Interface using Blocks | |
| with gr.Blocks(title="Neural Image Captioning", theme=gr.themes.Soft()) as interface: | |
| gr.Markdown("# 🖼️ Neural Image Captioning") | |
| gr.Markdown("Upload an image to generate a description using a CNN-LSTM model trained on Flickr8k.") | |
| status_text = gr.Markdown(load_model()) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| image_input = gr.Image(type="pil", label="Upload Image") | |
| with gr.Accordion("Advanced Settings", open=False): | |
| strategy_input = gr.Radio(["Beam Search", "Greedy"], value="Beam Search", label="Decoding Strategy") | |
| beam_size_input = gr.Slider(1, 10, value=3, step=1, label="Beam Size (for Beam Search)") | |
| max_len_input = gr.Slider(5, 50, value=20, step=1, label="Max Caption Length") | |
| temp_input = gr.Slider(0.1, 2.0, value=1.0, step=0.1, label="Temperature") | |
| submit_btn = gr.Button("Generate Caption", variant="primary") | |
| with gr.Column(scale=1): | |
| output_text = gr.Textbox(label="Generated Output", lines=4) | |
| submit_btn.click( | |
| fn=generate_caption, | |
| inputs=[image_input, strategy_input, beam_size_input, max_len_input, temp_input], | |
| outputs=output_text | |
| ) | |
| if __name__ == "__main__": | |
| # Run the Gradio app (defaults to port 7860, binds to 127.0.0.1 locally so the link works on Windows) | |
| interface.launch(server_port=7860, share=False) | |