Update app.py

app.py

@@ -1,3 +1,37 @@
 import gradio as gr
+from transformers import AutoModelForImageClassification, AutoProcessor
+import torch
 
-gr.load("models/dima806/facial_age_image_detection").launch()
+# 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)