Big overhaul to app.py

app.py

@@ -1,28 +1,27 @@
-from flask import Flask, render_template, request, jsonify, send_from_directory, send_file, Response
-from multiprocessing import Pool, cpu_count
-from threading import Thread
-from pathlib import Path
-from PIL import Image
-from datetime import datetime
 import os
-import tempfile
 import uuid
-import pandas as pd
-from werkzeug.utils import secure_filename
 import traceback
 import sys
 import io
 import zipfile
 import cv2
 import csv
-import numpy as np
-import redis
+import torch
 import json
 import shutil
-
+import numpy as np
+import pandas as pd
+from flask import Flask, Response, render_template, request, jsonify, send_from_directory, send_file, session, redirect, url_for
+from multiprocessing.pool import ThreadPool
+from threading import Thread
+from pathlib import Path
+from PIL import Image
+from datetime import datetime
+from werkzeug.utils import secure_filename
 from yolo_utils import load_model, detect_image
 
 app = Flask(__name__)
+app.secret_key = os.environ.get('FLASK_SECRET_KEY', str(uuid.uuid4()))  # For session security
 
 APP_ROOT = Path(__file__).parent
 UPLOAD_FOLDER = APP_ROOT / 'uploads'
@@ -35,9 +34,6 @@ app.config['ALLOWED_EXTENSIONS'] = {'png', 'jpg', 'jpeg', 'tif', 'tiff'}
 UPLOAD_FOLDER.mkdir(parents=True, exist_ok=True)
 RESULT_FOLDER.mkdir(parents=True, exist_ok=True)
 
-# Redis client (localhost:6379, db=0, no password)
-redis_client = redis.Redis(host='localhost', port=6379, db=0, decode_responses=True)
-
 @app.errorhandler(Exception)
 def handle_exception(e):
     print(f"Unhandled exception: {str(e)}")
@@ -51,73 +47,90 @@ def allowed_file(filename):
 def index():
     return render_template('index.html')
 
-# Global model for each process
+# Load model once at startup, use CUDA if available
+MODEL_DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 _model = None
 def get_model():
     global _model
     if _model is None:
         _model = load_model(WEIGHTS_FILE)
+        if MODEL_DEVICE == 'cuda':
+            _model.to('cuda')
     return _model
 
-def cleanup_job(job_id):
-    # Remove files
-    upload_dir = os.path.join(app.config['UPLOAD_FOLDER'], job_id)
-    if os.path.exists(upload_dir):
+def cleanup_session(session_id):
+    # Remove files for this session
+    upload_dir = Path(app.config['UPLOAD_FOLDER']) / session_id
+    result_dir = Path(app.config['RESULT_FOLDER']) / session_id
+    for d in [upload_dir, result_dir]:
+        if d.exists():
+            shutil.rmtree(d)
+
+# save the uploaded files
+@app.route('/uploads/<session_id>', methods=['POST'])
+def upload_files(session_id):
+    files = request.files.getlist('files')
+    upload_dir = Path(app.config['UPLOAD_FOLDER']) / session_id
+    # clear out any existing files for the session
+    if upload_dir.exists():
         shutil.rmtree(upload_dir)
-    # Remove Redis state
-    redis_client.delete(f"job:{job_id}")
-
-@app.route('/cleanup/<job_id>', methods=['POST'])
-def cleanup_job_endpoint(job_id):
-    cleanup_job(job_id)
-    return jsonify({'status': 'cleaned'})
-
-def get_job_state(job_id):
-    data = redis_client.get(f"job:{job_id}")
-    return json.loads(data) if data else None
-
-def set_job_state(job_id, state):
-    redis_client.set(f"job:{job_id}", json.dumps(state))
+    upload_dir.mkdir(parents=True, exist_ok=True)
+    filename_map = {}
+    file_data = []
+    for f in files:
+        orig_name = secure_filename(f.filename)
+        ext = Path(orig_name).suffix
+        unique_name = f"{uuid.uuid4().hex}{ext}"
+        file_path = upload_dir / unique_name
+        f.save(str(file_path))
+        filename_map[orig_name] = unique_name
+        with open(file_path, 'rb') as imgf:
+            file_data.append((orig_name, unique_name, imgf.read()))
+    return filename_map, file_data
 
-all_detections = {}
 
 def process_image(args):
     orig_name, unique_name, image_bytes = args
     model = get_model()
     detections = detect_image(model, image_bytes, conf=0.05)
-    # Save original image to uploads for later annotation (already saved)
     return {'orig_name': orig_name, 'unique_name': unique_name, 'detections': detections}
 
-def async_process_images(job_id, file_data):
+def async_process_images(session_id, file_data, state):
     try:
-        job_state = get_job_state(job_id)
-        job_state['status'] = 'running'
-        job_state['progress'] = 0
-        set_job_state(job_id, job_state)
+        state['status'] = 'running'
+        state['progress'] = 0
         total = len(file_data)
         results = []
         detections = {}
-        with Pool(processes=min(cpu_count(), total)) as pool:
-            for idx, result in enumerate(pool.imap(process_image, file_data)):
+        # Use ThreadPool for CPU, else single-threaded for CUDA
+        if MODEL_DEVICE == 'cuda':
+            pool = None
+            for idx, args in enumerate(file_data):
+                result = process_image(args)
                 results.append({
                     'filename': result['orig_name'],
                     'num_eggs': sum(1 for d in result['detections'] if d.get('class') == 'egg'),
                 })
                 detections[result['orig_name']] = result['detections']
-                # Update progress
-                job_state['progress'] = int((idx + 1) / total * 100)
-                set_job_state(job_id, job_state)
-        job_state['status'] = 'success'
-        job_state['results'] = results
-        job_state['detections'] = detections
-        job_state['progress'] = 100
-        set_job_state(job_id, job_state)
+                state['progress'] = int((idx + 1) / total * 100)
+        else:
+            with ThreadPool() as pool:
+                for idx, result in enumerate(pool.imap(process_image, file_data)):
+                    results.append({
+                        'filename': result['orig_name'],
+                        'num_eggs': sum(1 for d in result['detections'] if d.get('class') == 'egg'),
+                    })
+                    detections[result['orig_name']] = result['detections']
+                    state['progress'] = int((idx + 1) / total * 100)
+        state['status'] = 'success'
+        state['results'] = results
+        state['detections'] = detections
+        state['progress'] = 100
     except Exception as e:
-        job_state = get_job_state(job_id) or {}
-        job_state['status'] = 'error'
-        job_state['error'] = str(e)
-        job_state['progress'] = 100
-        set_job_state(job_id, job_state)
+        state['status'] = 'error'
+        state['error'] = str(e)
+        state['progress'] = 100
+
 
 @app.route('/process', methods=['POST'])
 def process_images():
@@ -125,80 +138,62 @@ def process_images():
         files = request.files.getlist('files')
         if not files or files[0].filename == '':
             return jsonify({'error': 'No files uploaded'}), 400
-        job_id = str(uuid.uuid4())
-        # Clean up any previous state for this job
-        cleanup_job(job_id)
-        filename_map, file_data = save_uploaded_files(files, job_id)
-        # Store initial job state in Redis
-        job_state = {
+        # Assign a session ID if not present
+        if 'id' not in session:
+            session['id'] = str(uuid.uuid4())
+        session_id = session['id']
+        # Clean up any previous state for this session
+        cleanup_session(session_id)
+        filename_map, file_data = upload_files(files, session_id)
+        # Store job state in session
+        state = {
             'status': 'starting',
             'progress': 0,
             'results': [],
             'filename_map': filename_map,
             'detections': {},
         }
-        set_job_state(job_id, job_state)
-        thread = Thread(target=async_process_images, args=(job_id, file_data))
+        session['job_state'] = state
+        thread = Thread(target=async_process_images, args=(session_id, file_data, state))
         thread.daemon = True
         thread.start()
-        return jsonify({'jobId': job_id})
+        return jsonify({'jobId': session_id})
     except Exception as e:
         print(f"Error in /process: {e}")
         print(traceback.format_exc())
         return jsonify({'error': str(e)}), 500
 
-def save_uploaded_files(files, job_id):
-    upload_dir = os.path.join(app.config['UPLOAD_FOLDER'], job_id)
-    if os.path.exists(upload_dir):
-        shutil.rmtree(upload_dir)
-    os.makedirs(upload_dir, exist_ok=True)
-    filename_map = {}
-    file_data = []
-    for f in files:
-        orig_name = secure_filename(f.filename)
-        ext = os.path.splitext(orig_name)[1]
-        unique_name = f"{uuid.uuid4().hex}{ext}"
-        file_path = os.path.join(upload_dir, unique_name)
-        f.save(file_path)
-        filename_map[orig_name] = unique_name
-        with open(file_path, 'rb') as imgf:
-            file_data.append((orig_name, unique_name, imgf.read()))
-    return filename_map, file_data
-
-@app.route('/progress/<job_id>')
-def get_progress(job_id):
-    job_state = get_job_state(job_id)
+@app.route('/progress/<session_id>')
+def get_progress(session_id):
+    # Only allow access to own session
+    if 'id' not in session or session['id'] != session_id:
+        return jsonify({"status": "error", "error": "Session not found or expired"}), 404
+    job_state = session.get('job_state')
     if not job_state:
-        return jsonify({"status": "error", "error": "Job ID not found"}), 404
-    # Add a mapping from filename to detections for frontend plotting
+        return jsonify({"status": "error", "error": "No job state"}), 404
     if 'detections' in job_state:
         job_state['detections_by_filename'] = job_state['detections']
     return jsonify(job_state)
 
-@app.route('/results/<job_id>/<path:filename>')
-def download_file(job_id, filename):
+@app.route('/results/<session_id>/<path:filename>')
+def download_file(session_id, filename):
     try:
-        try:
-            uuid.UUID(job_id, version=4)
-        except ValueError:
-            return jsonify({"error": "Invalid job ID format"}), 400
 
+        if 'id' not in session or session['id'] != session_id:
+            return jsonify({"error": "Session not found or expired"}), 404
         if '..' in filename or filename.startswith('/'):
             return jsonify({"error": "Invalid filename"}), 400
-
         safe_filename = secure_filename(filename)
-        file_dir = Path(app.config['RESULT_FOLDER']) / job_id
+        file_dir = Path(app.config['RESULT_FOLDER']) / session_id
         file_path = (file_dir / safe_filename).resolve()
-
         if not str(file_path).startswith(str(file_dir.resolve())):
-            print(f"Attempted path traversal: {job_id}/{filename}")
+            print(f"Attempted path traversal: {session_id}/{filename}")
             return jsonify({"error": "Invalid file path"}), 400
-
         if not file_path.is_file():
             if not file_dir.exists():
-                return jsonify({"error": f"Job directory {job_id} not found"}), 404
+                return jsonify({"error": f"Session directory {session_id} not found"}), 404
             files_in_dir = list(file_dir.iterdir())
-            return jsonify({"error": f"File '{filename}' not found in job '{job_id}'. Available: {[f.name for f in files_in_dir]}"}), 404
+            return jsonify({"error": f"File '{filename}' not found in session '{session_id}'. Available: {[f.name for f in files_in_dir]}"}), 404
 
         if filename.lower().endswith(('.tif', '.tiff')):
             try:
@@ -240,17 +235,15 @@ def download_file(job_id, filename):
         print(error_message)
         return jsonify({"error": "Server error", "log": error_message}), 500
 
-@app.route('/export_images/<job_id>')
-def export_images(job_id):
+@app.route('/export_images/<session_id>')
+def export_images(session_id):
     try:
-        try:
-            uuid.UUID(job_id, version=4)
-        except ValueError:
-            return jsonify({"error": "Invalid job ID format"}), 400
 
-        job_dir = Path(app.config['RESULT_FOLDER']) / job_id
+        if 'id' not in session or session['id'] != session_id:
+            return jsonify({"error": "Session not found or expired"}), 404
+        job_dir = Path(app.config['RESULT_FOLDER']) / session_id
         if not job_dir.exists():
-            return jsonify({"error": f"Job directory {job_id} not found"}), 404
+            return jsonify({"error": f"Session directory {session_id} not found"}), 404
 
         annotated_files = list(job_dir.glob('*_annotated.*'))
         if not annotated_files:
@@ -280,9 +273,9 @@ def export_images(job_id):
 def export_csv():
     try:
         data = request.json
-        job_id = data['jobId']
+        session_id = session.get('id')
         threshold = float(data.get('confidence', 0.5))
-        job_state = get_job_state(job_id)
+        job_state = session.get('job_state')
         if not job_state:
             return jsonify({'error': 'Job not found'}), 404
         rows = []
@@ -311,17 +304,17 @@ def export_csv():
 def export_images_post():
     try:
         data = request.json
-        job_id = data['jobId']
+        session_id = session.get('id')
         threshold = float(data.get('confidence', 0.5))
-        job_state = get_job_state(job_id)
+        job_state = session.get('job_state')
         if not job_state:
             return jsonify({'error': 'Job not found'}), 404
         memory_file = io.BytesIO()
         with zipfile.ZipFile(memory_file, 'w', zipfile.ZIP_DEFLATED) as zf:
             for orig_name, detections in job_state['detections'].items():
                 unique_name = job_state['filename_map'][orig_name]
-                img_path = os.path.join(app.config['UPLOAD_FOLDER'], job_id, unique_name)
-                img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
+                img_path = Path(app.config['UPLOAD_FOLDER']) / session_id / unique_name
+                img = cv2.imread(str(img_path), cv2.IMREAD_UNCHANGED)
                 filtered = [d for d in detections if d['score'] >= threshold]
                 for det in filtered:
                     x1, y1, x2, y2 = map(int, det['bbox'])
@@ -356,7 +349,7 @@ def print_startup_info():
         except Exception as e:
             print(f"Could not get weights file size: {e}")
 
-    is_container = os.path.exists('/.dockerenv') or 'DOCKER_HOST' in os.environ
+    is_container = Path('/.dockerenv').exists() or 'DOCKER_HOST' in os.environ
     print(f"Running in container: {is_container}")
     
     if is_container:
@@ -386,6 +379,17 @@ def print_startup_info():
         except Exception as e:
             print(f"Could not get NemaQuant script details: {e}")
 
+@app.before_request
+def ensure_session_id():
+    if 'id' not in session:
+        session['id'] = str(uuid.uuid4())
+
+@app.teardown_appcontext
+def cleanup_on_teardown(exception):
+    # If session is gone, clean up files
+    if 'id' in session and not session.modified:
+        cleanup_session(session['id'])
+
 if __name__ == '__main__':
     print_startup_info()
     app.run(host='0.0.0.0', port=7860, debug=True)