Team Ai
Apppublic

karan99300/ImageClassification

sourceHugging Faceupdated 2y agoView on Hugging Face
0likes
app.py51 linesDownload Raw Back to root
1from flask import Flask, request, jsonify, render_template2from transformers import AutoFeatureExtractor, AutoModelForImageClassification3from PIL import Image4import requests5import torch6 7# Initialize Flask app8app = Flask(__name__)9 10# Load pre-trained model and feature extractor11feature_extractor = AutoFeatureExtractor.from_pretrained('karan99300/ConvNext-finetuned-CIFAR100')12model = AutoModelForImageClassification.from_pretrained('karan99300/ConvNext-finetuned-CIFAR100')13 14# Define route for home page with form15@app.route('/', methods=['GET', 'POST'])16def index():17    if request.method == 'POST':18        # Get image URL from form submission19        image_url = request.form['image_url']20        21        # Classify image22        predicted_class = classify_image(image_url)23        24        return render_template('index.html', predicted_class=predicted_class, image_url=image_url)25    26    return render_template('index.html')27 28# Function to classify image29def classify_image(image_url):30    # Fetch image from URL31    try:32        image = Image.open(requests.get(image_url, stream=True).raw)33    except Exception as e:34        return f'Error fetching image: {str(e)}'35    36    # Preprocess image and perform inference37    pixel_values = feature_extractor(image.convert('RGB'), return_tensors='pt').pixel_values38    with torch.no_grad():39        outputs = model(pixel_values)40        logits = outputs.logits41        predicted_class_idx = logits.argmax(-1).item()42    43    # Get predicted label44    predicted_label = model.config.id2label[predicted_class_idx]45    46    return predicted_label47 48# Run Flask app49if __name__ == '__main__':50    app.run(debug=True,port=5000)51