File size: 1,756 Bytes
58a226a
 
bd0320b
58a226a
 
61a66af
 
58a226a
bd0320b
58a226a
61a66af
8c9f732
bd0320b
58a226a
 
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('static')  # Cache local des modèles
    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)