adding ui

app.py

@@ -1,6 +1,7 @@
 from transformers import AutoModelForImageClassification, AutoFeatureExtractor
 import torch
 from PIL import Image
+import gradio as gr
 
 def load_model_from_hub(repo_id: str):
     """
@@ -39,26 +40,54 @@ def predict(image_path: str, model, processor):
     
     return predictions
 
-# Example usage in your Flask/FastAPI app
-from flask import Flask, request
-app = Flask(__name__)
+def predict_image(image):
+    """
+    Gradio interface function for prediction
+    
+    Args:
+        image: Image uploaded through Gradio interface
+    Returns:
+        str: Prediction result with confidence score
+    """
+    # Convert from numpy array to PIL Image
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+    
+    # Process image and get prediction
+    inputs = processor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+        predictions = outputs.logits.softmax(-1)
+    
+    # Get the top prediction
+    pred_scores = predictions[0].tolist()
+    top_pred_idx = max(range(len(pred_scores)), key=pred_scores.__getitem__)
+    confidence = pred_scores[top_pred_idx]
+    
+    # Get class label
+    if hasattr(model.config, 'id2label'):
+        label = model.config.id2label[top_pred_idx]
+    else:
+        label = f"Class {top_pred_idx}"
+    
+    return f"{label} (Confidence: {confidence:.2%})"
 
 # Load model at startup
 model, processor = load_model_from_hub("srtangirala/resnet50-exp")
 
-@app.route('/predict', methods=['POST'])
-def predict_endpoint():
-    if 'file' not in request.files:
-        return {'error': 'No file provided'}, 400
-    
-    file = request.files['file']
-    image_path = "temp_image.jpg"  # You might want to generate a unique filename
-    file.save(image_path)
-    
-    predictions = predict(image_path, model, processor)
-    
-    # Convert predictions to list and return
-    return {'predictions': predictions.tolist()[0]}
+# Create Gradio interface
+iface = gr.Interface(
+    fn=predict_image,
+    inputs=gr.Image(),
+    outputs=gr.Text(),
+    title="Image Classification",
+    description="Upload an image to classify it!",
+    examples=[
+        # You can add example images here
+        # ["path/to/example1.jpg"],
+        # ["path/to/example2.jpg"]
+    ]
+)
 
-if __name__ == '__main__':
-    app.run(debug=True) 
+if __name__ == "__main__":
+    iface.launch()