Spaces:
Sleeping
Sleeping
import gradio as gr | |
from transformers import AutoModelForImageClassification, AutoProcessor | |
import torch | |
# Load the model and processor from Hugging Face | |
model_name = "dima806/facial_age_image_detection" | |
model = AutoModelForImageClassification.from_pretrained(model_name) | |
processor = AutoProcessor.from_pretrained(model_name) | |
# Define the prediction function | |
def predict(image): | |
# Process the input image | |
inputs = processor(images=image, return_tensors="pt") | |
# Perform the prediction | |
with torch.no_grad(): | |
outputs = model(**inputs) | |
# Get the model's original outputs (e.g., logits or probabilities) | |
predictions = outputs.logits | |
# Convert predictions to a list and round to 2 decimal places if necessary | |
predictions_list = predictions.tolist() | |
rounded_predictions = [[round(pred, 2) for pred in prediction] for prediction in predictions_list] | |
return rounded_predictions | |
# Create Gradio interface | |
iface = gr.Interface( | |
fn=predict, | |
inputs="image", | |
outputs="label", # Use the model's original output type | |
title="Facial Age Prediction", | |
description="This application predicts your age from a facial image." | |
) | |
# Launch the Gradio application | |
iface.launch(share=True) | |