Update server.py

server.py

@@ -1,12 +1,14 @@
 import torch
 import torch.nn as nn
-from flask import Flask, request, jsonify, render_template
+from flask import Flask, request, jsonify, render_template, make_response
 from flask_cors import CORS
 import io
 import os
 from PIL import Image
 from diffusers import StableDiffusionPipeline
 import os
+from tinydb import TinyDB, Query
+import uuid
 
 token = os.getenv("HF_TOKEN")
 
@@ -33,6 +35,8 @@ CORS(app)
 stable_diff_pipe = None
 model = None
 
+db = TinyDB('artwork_results.json')
+
 def load_models(model_name="rupeshs/LCM-runwayml-stable-diffusion-v1-5"):
     global stable_diff_pipe, model
     
@@ -60,6 +64,13 @@ def extract_image_features(image):
 
     return generated_features
 
+# Helper to get or set a unique user_id via cookies
+def get_user_id():
+    user_id = request.cookies.get('user_id')
+    if not user_id:
+        user_id = str(uuid.uuid4())
+    return user_id
+
 @app.route('/')
 def index():
     return render_template('index.html')
@@ -67,45 +78,58 @@ def index():
 @app.route('/api/check-membership', methods=['POST'])
 def check_membership():
     try:
-        # Get the model name from the request
         model_name = request.form.get('model', 'rupeshs/LCM-runwayml-stable-diffusion-v1-5')
-        
+        contact_info = request.form.get('contact_info')  # Optional, if frontend sends it
+        opt_in = request.form.get('opt_in', 'false').lower() == 'true'
+        user_id = get_user_id()
         # Ensure models are loaded with the selected model
+        global stable_diff_pipe, model
         if stable_diff_pipe is None or model is None:
             load_models(model_name)
         elif stable_diff_pipe.name_or_path != model_name:
-            # Reload the model if a different one is selected
             load_models(model_name)
-            
         if 'image' not in request.files:
             return jsonify({'error': 'No image found in request'}), 400
-        
-        # Get the image from the request
         file = request.files['image']
         image_bytes = file.read()
         image = Image.open(io.BytesIO(image_bytes))
-
-        # Get image features using Stable Diffusion
         image_features = extract_image_features(image)
-        
-        # Preprocess the features for MIDM model
-        processed_features = image_features.reshape(1, -1)[:, :10]  # Select first 10 features
-        
-        # Perform inference
+        processed_features = image_features.reshape(1, -1)[:, :10]
         with torch.no_grad():
             output = model(processed_features)
             probability = output.item()
             predicted = int(output > 0.5)
-            
-        return jsonify({
+        # Store in DB if likely and opt-in
+        if predicted == 1:
+            entry = {
+                'model': model_name,
+                'user_id': user_id,
+                'timestamp': int(uuid.uuid1().time),
+            }
+            if opt_in and contact_info:
+                entry['contact_info'] = contact_info
+            # Only store if not already present for this user/model
+            User = Query()
+            exists = db.search((User.model == model_name) & (User.user_id == user_id))
+            if not exists:
+                db.insert(entry)
+        # Count unique user_ids for this model
+        User = Query()
+        user_ids = set(r['user_id'] for r in db.search(User.model == model_name))
+        count = len(user_ids)
+        response = make_response(jsonify({
             'probability': probability,
             'predicted_class': predicted,
             'message': f"Predicted membership probability: {probability}",
-            'is_in_training_data': "Likely" if predicted == 1 else "Unlikely"
-        })
-        
+            'is_in_training_data': "Likely" if predicted == 1 else "Unlikely",
+            'likely_count': count
+        }))
+        # Set user_id cookie if not present
+        if not request.cookies.get('user_id'):
+            response.set_cookie('user_id', user_id, max_age=60*60*24*365*5)  # 5 years
+        return response
     except Exception as e:
-        print(f"Error processing request: {str(e)}")  # Add logging
+        print(f"Error processing request: {str(e)}")
         return jsonify({'error': str(e)}), 500
 
 if __name__ == '__main__':