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)