Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import pipeline | |
| import torch | |
| import os | |
| from PIL import Image | |
| import io | |
| # Check if CUDA is available (Hugging Face Spaces supports GPU acceleration) | |
| device = 0 if torch.cuda.is_available() else -1 | |
| print(f"Using device: {'CUDA' if device == 0 else 'CPU'}") | |
| # Initialize the image processing pipeline | |
| # You can replace this with any Hugging Face model that processes images | |
| model_name = "google/vit-base-patch16-224" | |
| image_processor = pipeline("image-classification", model=model_name, device=device) | |
| def process_image(input_image): | |
| """ | |
| Process the uploaded image through the model and return results | |
| """ | |
| if input_image is None: | |
| return [{"label": "No image provided", "score": 0.0}] | |
| # Run the image through the model | |
| results = image_processor(input_image) | |
| # Return top 5 predictions | |
| return results[:5] | |
| def save_output(results): | |
| """ | |
| Convert results to a downloadable format | |
| """ | |
| if not results or len(results) == 0: | |
| return None | |
| output_text = "Model Predictions:\n\n" | |
| for result in results: | |
| output_text += f"Label: {result['label']}, Score: {result['score']:.4f}\n" | |
| # Create a file for download | |
| with open("results.txt", "w") as f: | |
| f.write(output_text) | |
| return "results.txt" | |
| # Create the Gradio interface with a more polished design | |
| with gr.Blocks(theme=gr.themes.Soft()) as demo: | |
| gr.Markdown("# Image Classification Demo") | |
| gr.Markdown("Upload an image and get classification results from the model.") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| # Input components | |
| input_image = gr.Image(type="pil", label="Upload Image") | |
| with gr.Row(): | |
| submit_btn = gr.Button("Process Image", variant="primary") | |
| clear_btn = gr.Button("Clear") | |
| with gr.Column(scale=1): | |
| # Output components | |
| output_results = gr.JSON(label="Model Predictions") | |
| download_btn = gr.Button("Download Results") | |
| download_output = gr.File(label="Download Output") | |
| # Set up the processing flow | |
| submit_btn.click( | |
| fn=process_image, | |
| inputs=[input_image], | |
| outputs=[output_results] | |
| ) | |
| clear_btn.click( | |
| fn=lambda: (None, None, None), | |
| inputs=[], | |
| outputs=[input_image, output_results, download_output] | |
| ) | |
| download_btn.click( | |
| fn=save_output, | |
| inputs=[output_results], | |
| outputs=[download_output] | |
| ) | |
| # Add example images | |
| gr.Examples( | |
| examples=[ | |
| os.path.join(os.path.dirname(__file__), "examples/cat.jpg"), | |
| os.path.join(os.path.dirname(__file__), "examples/dog.jpg"), | |
| ], | |
| inputs=input_image, | |
| label="Example Images" | |
| ) | |
| # Add information footer | |
| gr.Markdown(""" | |
| ### About this demo | |
| - Model: [google/vit-base-patch16-224](https://huggingface.co/google/vit-base-patch16-224) | |
| - This demo classifies images into 1000 ImageNet categories | |
| - Created with Gradio and Hugging Face Transformers | |
| """) | |
| # For Hugging Face Spaces, we use the Gradio app directly | |
| demo.launch() |