Spaces:
Running
Running
File size: 1,770 Bytes
58a226a bd0320b 58a226a 61a66af 58a226a bd0320b 58a226a 61a66af 8c9f732 bd0320b 58a226a bd0320b e8c243a bd0320b 58a226a bd0320b 58a226a 4f5a1cf f5ae429 58a226a bd0320b 4f5a1cf bd0320b 4f5a1cf 58a226a bd0320b 58a226a bd0320b 58a226a 61a66af |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
# docker build -t reward-simulator .docker run -p 7860:7860 -v $(pwd)/data:/app/data reward-simulator
from flask import Flask, request, jsonify, render_template, send_from_directory
from PIL import Image
import numpy as np
import io
import torch
from request import get_ft # get_ft(model, image) doit retourner un np.ndarray
app = Flask(__name__)
# Global model
model = None
def load_model():
"""Load DINOv2 model"""
torch.hub.set_dir("/tmp/torch_cache") # Dossier temporaire autorisé
model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14')
model.eval()
model.to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))
return model
def init_model():
global model
model = load_model()
@app.route('/')
def home():
return render_template('index.html') # Si tu as un front-end intégré
@app.route('/static/<path:filename>')
def serve_static(filename):
return send_from_directory('static', filename)
@app.route('/process', methods=['POST'])
def process_images():
if 'image1' not in request.files or 'image2' not in request.files:
return jsonify({'error': 'Two images must be provided (image1 and image2)'}), 400
try:
image1 = Image.open(io.BytesIO(request.files['image1'].read())).convert('RGB')
image2 = Image.open(io.BytesIO(request.files['image2'].read())).convert('RGB')
features1 = get_ft(model, image1)
features2 = get_ft(model, image2)
distance = float(np.linalg.norm(features1 - features2))
return jsonify({'distance': distance})
except Exception as e:
print(f"Erreur back-end: {e}")
return jsonify({'error': str(e)}), 500
if __name__ == '__main__':
init_model()
app.run(host='0.0.0.0', port=7860)
|