imagetotextserver.py

#!/usr/bin/env python3

from flask import Flask, request, jsonify
from PIL import Image
from transformers import Blip2Processor, Blip2ForConditionalGeneration
import torch
import io

app = Flask(__name__)

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Device: {device}")

print("Loading processor...")
processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
print("Loading model...")
model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b")
print("Moving model to device...")
model.to(device)
print("Model loaded")

@app.route('/process', methods=['POST'])
def process_images():
    if 'images' not in request.files:
        return jsonify({"error": "No images part"}), 400
    
    files = request.files.getlist('images')
    images = []

    for file in files:
        try:
            image = Image.open(io.BytesIO(file.read()))
            images.append(image)
        except IOError:
            return jsonify({"error": "Invalid image"}), 400

    inputs = processor(images=images, return_tensors="pt").to(device)
    generated_ids = model.generate(**inputs)
    generated_texts = processor.batch_decode(generated_ids, skip_special_tokens=True)

    response = []
    for image, caption in zip(files, generated_texts):
        result = {
            "image_name": image.filename,
            "image_caption_result": caption.strip()
        }
        response.append(result)

    return jsonify(response)

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0')