Spaces:
Sleeping
Sleeping
File size: 3,234 Bytes
180415e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
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() |