app.py

import os
import uuid
import json
import traceback
import sys
import time
import io
import zipfile
import cv2
import csv
import pickle
import shutil
import logging
from ultralytics import YOLO
import numpy as np
import pandas as pd
from torch import cuda
from flask import Flask, Response, render_template, request, jsonify, send_file, session
from multiprocessing.pool import Pool
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from PIL import Image
from datetime import datetime
from werkzeug.utils import secure_filename
from yolo_utils import detect_in_image

app = Flask(__name__)
_secret_key = os.environ.get('FLASK_SECRET_KEY')
if not _secret_key:
    # Fallback for local dev only — sessions won't persist across restarts.
    # On HF Spaces, set FLASK_SECRET_KEY as a Space secret to avoid session loss between workers.
    _secret_key = str(uuid.uuid4())
    print("WARNING: FLASK_SECRET_KEY not set — using random key. Sessions will break across workers/restarts.")
else:
    print(f"INFO: FLASK_SECRET_KEY is set (length={len(_secret_key)})")
app.secret_key = _secret_key

# HF Spaces serves over HTTPS via a reverse proxy and may embed the app in an iframe.
# SameSite=None;Secure is required so cookies are sent in cross-site/iframe POST requests.
# HF sets SPACE_HOST env var; fall back to checking SPACE_ID or SPACE_AUTHOR_NAME.
_on_https = any(os.environ.get(v) for v in ('SPACE_HOST', 'SPACE_ID', 'SPACE_AUTHOR_NAME'))
app.config['SESSION_COOKIE_SECURE'] = _on_https
app.config['SESSION_COOKIE_SAMESITE'] = 'None' if _on_https else 'Lax'
app.config['SESSION_COOKIE_HTTPONLY'] = True
print(f"INFO: SESSION_COOKIE_SECURE={_on_https}, SAMESITE={'None' if _on_https else 'Lax'}")
# comment out these lines if you want to see full logs
log = logging.getLogger('werkzeug')
log.setLevel(logging.ERROR)

APP_ROOT = Path(__file__).parent
# Use /tmp for runtime data — works reliably on all container platforms.
# /tmp is RAM-backed tmpfs, always writable, avoids overlay filesystem issues on HF Spaces.
# Note: /tmp is cleared on container restart (uploads/results are transient by design).
UPLOAD_FOLDER = Path('/tmp/nemaquant/uploads')
RESULTS_FOLDER = Path('/tmp/nemaquant/results')
ANNOT_FOLDER = Path('/tmp/nemaquant/annotated')
SESSION_META_FOLDER = Path('/tmp/nemaquant/sessions')
WEIGHTS_FILE = APP_ROOT / 'weights.pt'
app.config['UPLOAD_FOLDER'] = str(UPLOAD_FOLDER)
app.config['RESULTS_FOLDER'] = str(RESULTS_FOLDER)
app.config['ANNOT_FOLDER'] = str(ANNOT_FOLDER)
app.config['WEIGHTS_FILE'] = str(WEIGHTS_FILE)
app.config['ALLOWED_EXTENSIONS'] = {'png', 'jpg', 'jpeg', 'tif', 'tiff'}

# Create dirs at startup
UPLOAD_FOLDER.mkdir(parents=True, exist_ok=True)
RESULTS_FOLDER.mkdir(parents=True, exist_ok=True)
ANNOT_FOLDER.mkdir(parents=True, exist_ok=True)
SESSION_META_FOLDER.mkdir(parents=True, exist_ok=True)
# YOLO_CONFIG_DIR points to /tmp/nemaquant/.yolo_config (set in Dockerfile ENV).
# Create it here so ultralytics can write its cache on read-only container filesystems
# (e.g. Apptainer SIF images).
Path(os.environ.get('YOLO_CONFIG_DIR', '/tmp/nemaquant/.yolo_config')).mkdir(parents=True, exist_ok=True)
print(f"Data root: /tmp/nemaquant | Weights: {WEIGHTS_FILE}")

# ---------------------------------------------------------------------------
# Session metadata helpers
# Flask's client-side cookie is limited to ~4KB. When many images are
# uploaded, filename_map / uuid_map_to_uuid_imgname can overflow.
# We persist them to disk so every route can recover them even when the
# cookie is absent or truncated (e.g. large batches, Apptainer --cleanenv,
# multi-worker gunicorn).
# ---------------------------------------------------------------------------
def _save_session_meta(session_id, filename_map, uuid_map):
    meta_dir = SESSION_META_FOLDER / session_id
    meta_dir.mkdir(parents=True, exist_ok=True)
    with open(meta_dir / 'meta.json', 'w') as fh:
        json.dump({'filename_map': filename_map, 'uuid_map_to_uuid_imgname': uuid_map}, fh)

def _load_session_meta(session_id):
    meta_path = SESSION_META_FOLDER / session_id / 'meta.json'
    if meta_path.exists():
        with open(meta_path) as fh:
            return json.load(fh)
    return {}

# Load model once at startup, use CUDA if available
MODEL_DEVICE = 'cuda' if cuda.is_available() else 'cpu'

# For GPU: load the model globally at startup so threads can reuse it without
# re-initialising CUDA (forked Pool workers cannot re-init CUDA in the child).
# For CPU: model is loaded per-worker in init_worker() instead.
_gpu_model = None
if MODEL_DEVICE == 'cuda':
    _gpu_model = YOLO(str(WEIGHTS_FILE))
    _gpu_model.to('cuda')
    print(f'GPU model loaded at startup on {MODEL_DEVICE}')

# Wrapper so GPU futures (concurrent.futures.Future) expose the same
# .ready() interface as multiprocessing AsyncResult.
class _FutureWrapper:
    def __init__(self, future):
        self._f = future
    def ready(self):
        return self._f.done()
    def get(self):
        return self._f.result()

# Global dict mapping session_id -> async result (Pool AsyncResult or _FutureWrapper)
async_results = {}

@app.errorhandler(Exception)
def handle_exception(e):
    print(f"Unhandled exception: {str(e)}")
    print(traceback.format_exc())
    return jsonify({"error": "Server error", "log": str(e)}), 500

@app.route('/')
def index():
    return render_template('index.html')

# save the uploaded files
@app.route('/uploads', methods=['POST'])
def upload_files():
    session_id = 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)
    upload_dir.mkdir(parents=True, exist_ok=True)
    # generate new unique filenames via uuid, save the mapping dict of old:new to session
    filename_map = {}
    uuid_map_to_uuid_imgname = {}
    for f in files:
        orig_name = secure_filename(f.filename)
        ext = Path(orig_name).suffix
        uuid_base = uuid.uuid4().hex
        uuid_name = f"{uuid_base}{ext}"
        file_path = upload_dir / uuid_name
        f.save(str(file_path))
        filename_map[uuid_base] = orig_name
        uuid_map_to_uuid_imgname[uuid_base] = uuid_name
    session['filename_map'] = filename_map
    session['uuid_map_to_uuid_imgname'] = uuid_map_to_uuid_imgname
    # Persist to disk — cookie may be silently dropped if it exceeds ~4KB
    _save_session_meta(session_id, filename_map, uuid_map_to_uuid_imgname)
    return jsonify({'filename_map': filename_map, 'session_id': session_id, 'status': 'uploaded'})

# /preview route for serving original uploaded image
@app.route('/preview', methods=['POST'])
def preview_image():
    try:
        data = request.get_json()
        uuid = data.get('uuid')
        # Prefer client-supplied session_id (cookie may differ on HF Spaces HTTPS proxy)
        session_id = data.get('session_id') or session['id']
        _meta = _load_session_meta(session_id)
        uuid_map_to_uuid_imgname = session.get('uuid_map_to_uuid_imgname') or _meta.get('uuid_map_to_uuid_imgname', {})
        img_name = uuid_map_to_uuid_imgname.get(uuid)
        if not img_name:
            print(f"/preview: No img_name found for uuid {uuid}")
            return jsonify({'error': 'File not found'}), 404
        img_path = Path(app.config['UPLOAD_FOLDER']) / session_id / img_name
        if not img_path.exists():
            print(f"/preview: File does not exist at {img_path}")
            return jsonify({'error': 'File not found'}), 404
        # Determine MIME type
        ext = img_path.suffix.lower()
        if ext in ['.jpg', '.jpeg']:
            mimetype = 'image/jpeg'
        elif ext in ['.png']:
            mimetype = 'image/png'
        elif ext in ['.tif', '.tiff']:
            mimetype = 'image/tiff'
        else:
            mimetype = 'application/octet-stream'
        return send_file(
            str(img_path),
            mimetype=mimetype,
            as_attachment=False,
            download_name=img_name
        )
    except Exception as e:
        print(f"Error in /preview: {e}")
        return jsonify({'error': str(e)}), 500

# initializer for Pool to load model in each process
# each worker will have its own model instance (CPU only)
def init_worker(model_path):
    global model
    model = YOLO(model_path)

# CPU pool worker — uses per-worker model loaded by init_worker()
def process_single_image(img_path, results_dir):
    global model
    uuid_base = img_path.stem
    pickle_path = results_dir / f"{uuid_base}.pkl"
    results = detect_in_image(model, str(img_path))
    with open(pickle_path, 'wb') as pf:
        pickle.dump(results, pf)
    return uuid_base

# GPU thread worker — reuses the global _gpu_model loaded at startup
def process_single_image_thread(img_path, results_dir):
    global _gpu_model
    uuid_base = img_path.stem
    pickle_path = results_dir / f"{uuid_base}.pkl"
    results = detect_in_image(_gpu_model, str(img_path))
    with open(pickle_path, 'wb') as pf:
        pickle.dump(results, pf)
    return uuid_base

@app.route('/process', methods=['POST'])
def start_processing():
    session_id = session['id']
    # The client echoes back the session_id it received from /uploads.
    # On HF Spaces the session cookie can be missing on subsequent requests
    # (HTTPS proxy / SameSite), so we fall back to the client-supplied id
    # when the cookie-based id doesn't have an upload directory.
    client_session_id = request.form.get('session_id', '')
    if client_session_id:
        # Prefer the client-supplied id unconditionally — it's the authoritative
        # id from the /uploads call; the cookie may point to a different worker session.
        client_dir = Path(app.config['UPLOAD_FOLDER']) / client_session_id
        if client_dir.exists() or not (Path(app.config['UPLOAD_FOLDER']) / session_id).exists():
            session_id = client_session_id
            session['id'] = session_id
    job_state = {
        "status": "starting",
        "progress": 0,
        "sessionId": session_id
    }
    session['job_state'] = job_state
    upload_dir = Path(app.config['UPLOAD_FOLDER']) / session_id
    results_dir = Path(app.config['RESULTS_FOLDER']) / session_id

    try:
        # Fail fast with a clear message if the upload directory is missing
        if not upload_dir.exists():
            available = [d.name for d in Path(app.config['UPLOAD_FOLDER']).iterdir()] \
                if Path(app.config['UPLOAD_FOLDER']).exists() else []
            msg = (f"Upload directory not found: {upload_dir}. "
                   f"cookie_session={session['id']}, client_session={request.form.get('session_id','')}, "
                   f"available={available}")
            print(f"ERROR /process: {msg}")
            return jsonify({'error': msg}), 500

        # clean out old results if needed
        if results_dir.exists():
            shutil.rmtree(results_dir)
        results_dir.mkdir(parents=True)

        # set up iterable of uploaded files to process
        arg_list = [(x, results_dir) for x in list(upload_dir.iterdir())]

        if MODEL_DEVICE == 'cuda':
            # GPU: run in a single thread so CUDA is never re-initialised in a
            # forked subprocess (Pool uses fork by default, which breaks CUDA).
            def _gpu_task():
                for img_path, res_dir in arg_list:
                    process_single_image_thread(img_path, res_dir)
            executor = ThreadPoolExecutor(max_workers=1)
            future = executor.submit(_gpu_task)
            executor.shutdown(wait=False)
            async_results[session_id] = _FutureWrapper(future)
        else:
            n_proc = os.cpu_count()
            pool = Pool(processes=n_proc,
                      initializer=init_worker,
                      initargs=(str(WEIGHTS_FILE),))
            async_results[session_id] = pool.starmap_async(process_single_image, arg_list)
            pool.close()

        # Update job state after process launch
        job_state["status"] = "processing"
        session['job_state'] = job_state
        return jsonify({'status': 'processing',
                        'sessionId': session_id
                        })
    except Exception as e:
        print(f"Error in /process: {e}")
        print(traceback.format_exc())
        return jsonify({'error': str(e),
                        'status': 'unknown',
                        'sessionId': session_id}), 500
    
@app.route('/progress')
def get_progress():
        session_id = session['id']
        # Accept client-supplied session_id as fallback (cookie may be missing on HF Spaces)
        client_session_id = request.args.get('session_id', '')
        if client_session_id and session_id not in async_results and client_session_id in async_results:
            session_id = client_session_id
            session['id'] = session_id
        try:
            job_state = session.get('job_state')
            # If session lost job_state but we have an async_result, reconstruct from disk
            if not job_state and session_id in async_results:
                job_state = {'status': 'processing', 'progress': 0, 'sessionId': session_id}
            if not job_state:
                print("/progress: No job_state found in session.")
                return jsonify({"status": "error", "error": "No job state"}), 404

            results_dir = Path(app.config['RESULTS_FOLDER']) / session_id
            uploads_dir = Path(app.config['UPLOAD_FOLDER']) / session_id
            n_results = len(list(results_dir.glob('*.pkl')))
            n_uploads = len(list(uploads_dir.iterdir()))

            # If async_result is ready, verify completion and update job state
            async_result = async_results.get(session_id)
            if async_result and async_result.ready():
                if n_results == n_uploads:
                    job_state['status'] = 'completed'
                    job_state['progress'] = 100
                    session['job_state'] = job_state
                    _meta = _load_session_meta(session_id)
                    _filename_map = session.get('filename_map') or _meta.get('filename_map', {})
                    resp = {
                        'status': 'completed',
                        'progress': 100,
                        'filename_map': _filename_map,
                        'session_id': job_state.get('sessionId'),
                        'error': job_state.get('error'),
                    }
                    # Aggregate results into a single response object
                    all_results = {}
                    for pkl_file in results_dir.glob('*.pkl'):
                        uuid_base = pkl_file.stem
                        with open(pkl_file, 'rb') as pf:
                            all_results[uuid_base] = pickle.load(pf)
                    resp['results'] = all_results
                    print(f"Job executed successfully! {len(all_results)} results aggregated.")
                    return jsonify(resp)
                
            # If still processing, update progress
            if job_state.get('status') == 'processing':
                progress = int((n_results / n_uploads) * 100) if n_uploads > 0 else 0
                job_state['progress'] = progress
                session['job_state'] = job_state
                resp = {
                    'status': 'processing',
                    'progress': progress,
                    'sessionId': session_id,
                }
                return jsonify(resp)
            # Default response as a catchall
            resp = {
                'status': job_state.get('status', 'unknown'),
                'progress': job_state.get('progress', 0),
                'sessionId': job_state.get('session_id'),
                'error': job_state.get('error'),
            }
            return jsonify(resp)
        except Exception as e:
            print(f"Error in /progress: {e}")
            print(traceback.format_exc())
            return jsonify({"status": "error", "error": str(e)}), 500

def read_img_and_draw(img_path, detections, confidence):
    img = cv2.imread(str(img_path), cv2.IMREAD_UNCHANGED)
    filtered = [d for d in detections if d.get('score', 0) >= confidence]
    for det in filtered:
        x1, y1, x2, y2 = map(int, det['bbox'])
        cv2.rectangle(img, (x1, y1), (x2, y2), (0,0,255), 3)
    return img

# /annotate route for dynamic annotation
@app.route('/annotate', methods=['POST'])
def annotate_image():
    try:
        data = request.get_json()
        uuid = data.get('uuid')
        confidence = float(data.get('confidence', 0.5))
        # Prefer client-supplied session_id (cookie may differ on HF Spaces HTTPS proxy)
        session_id = data.get('session_id') or session['id']
        _meta = _load_session_meta(session_id)
        uuid_map_to_uuid_imgname = session.get('uuid_map_to_uuid_imgname') or _meta.get('uuid_map_to_uuid_imgname', {})
        img_name = uuid_map_to_uuid_imgname.get(uuid)
        orig_img_name = (session.get('filename_map') or _meta.get('filename_map', {})).get(uuid)

        if not img_name:
            return jsonify({'error': 'File not found'}), 404
        # Load detections from pickle
        result_path = Path(app.config['RESULTS_FOLDER']) / session_id / f"{uuid}.pkl"
        if not result_path.exists():
            return jsonify({'error': 'Results not found'}), 404
        with open(result_path, 'rb') as pf:
            detections = pickle.load(pf)

        img_path = Path(app.config['UPLOAD_FOLDER']) / session_id / img_name
        img = read_img_and_draw(img_path, detections, confidence)
        # Save annotated image out
        annot_dir = Path(app.config['ANNOT_FOLDER']) / session_id
        annot_dir.mkdir(parents=True, exist_ok=True)
        annot_imgname = f"{uuid}_annotated.png"
        annot_imgpath = str(annot_dir / annot_imgname)
        cv2.imwrite(annot_imgpath, img)

        # Serve image directly from disk
        return send_file(
            annot_imgpath,
            mimetype='image/png',
            as_attachment=False,
            download_name=annot_imgname
        )
    except Exception as e:
        print(f"Error in /annotate: {e}")
        return jsonify({'error': str(e)}), 500

@app.route('/export_images', methods=['POST'])
def export_images():
    try:
        data = request.get_json()
        confidence = float(data.get('confidence', 0.5))
        session_id = data.get('session_id') or session['id']
        _meta = _load_session_meta(session_id)
        filename_map = session.get('filename_map') or _meta.get('filename_map', {})
        uuid_map_to_uuid_imgname = session.get('uuid_map_to_uuid_imgname') or _meta.get('uuid_map_to_uuid_imgname', {})
        # ensure there's a landing spot
        annot_dir = Path(app.config['ANNOT_FOLDER']) / session_id
        annot_dir.mkdir(parents=True, exist_ok=True)

        # add all annotated files to zip
        memory_file = io.BytesIO()
        with zipfile.ZipFile(memory_file, 'w', zipfile.ZIP_DEFLATED) as zf:
            # iterate through all uuids
            for uuid in filename_map.keys():
                img_name = uuid_map_to_uuid_imgname.get(uuid)
                if not img_name:
                    continue
                img_path = Path(app.config['UPLOAD_FOLDER']) / session_id / img_name
                result_path = Path(app.config['RESULTS_FOLDER']) / session_id / f"{uuid}.pkl"
                if not result_path.exists():
                    return jsonify({'error': 'Results not found'}), 404
                if not img_path.exists():
                    return jsonify({'error': 'Image not found'}), 404
                with open(result_path, 'rb') as pf:
                    detections = pickle.load(pf)
                img = read_img_and_draw(img_path, detections, confidence)
                # clean the name
                orig_name = filename_map.get(uuid)
                annot_imgname = f"{str(Path(orig_name).stem)}_annotated.png"
                annot_imgpath = str(annot_dir / annot_imgname)
                cv2.imwrite(annot_imgpath, img)
                zf.write(annot_imgpath, annot_imgname)
        # timestamp for filename
        memory_file.seek(0)
        timestamp = datetime.now().strftime('%Y%m%d-%H%M%S')
        
        return send_file(
            memory_file,
            mimetype='application/zip',
            as_attachment=True,
            download_name=f'nemaquant_annotated_{timestamp}.zip'
        )

    except Exception as e:
        error_message = f"Error exporting images: {str(e)}"
        print(error_message)
        return jsonify({"error": "Server error", "log": error_message}), 500

@app.route('/export_csv', methods=['POST'])
def export_csv():
    try:
        data = request.json
        session_id = data.get('session_id') or session['id']
        job_state = session.get('job_state')
        _meta = _load_session_meta(session_id)
        filename_map = session.get('filename_map') or _meta.get('filename_map', {})
        threshold = float(data.get('confidence', 0.5))
        if not job_state:
            return jsonify({'error': 'Job not found'}), 404
        
        # iterate through the results
        results_dir = Path(app.config['RESULTS_FOLDER']) / session_id
        pkl_paths = list(results_dir.glob('*.pkl'))
        all_results = {}
        for path in pkl_paths:
            uuid_base = path.stem
            with open(path, 'rb') as pf:
                all_results[uuid_base] = pickle.load(pf)

        # populate rows for CSV conversion
        rows = []
        for uuid in all_results.keys():
            count = sum(1 for d in all_results[uuid] if d['score'] >= threshold)
            rows.append({'Filename': filename_map.get(uuid, uuid), 'EggsDetected': count, 'ConfidenceThreshold': threshold})
        rows = sorted(rows, key=lambda x: x['Filename'].lower())
        # write the CSV out
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        output = io.StringIO()
        writer = csv.DictWriter(output, fieldnames=['Filename', 'EggsDetected', 'ConfidenceThreshold'])
        writer.writeheader()
        writer.writerows(rows)
        output.seek(0)
        return Response(
            output.getvalue(),
            mimetype='text/csv',
            headers={
                'Content-Disposition': f'attachment; filename=nemaquant_results_{timestamp}.csv'
            }
        )
    except Exception as e:
        error_message = f"Error exporting CSV: {str(e)}"
        print(error_message)
        return jsonify({"error": "Server error", "log": error_message}), 500

@app.before_request
def ensure_session():
    if 'id' not in session:
        session['id'] = uuid.uuid4().hex


def print_startup_info():
    print("----- NemaQuant Flask App Starting -----")
    print(f"Working directory: {os.getcwd()}")
    python_version_single_line = sys.version.replace('\n', ' ')
    print(f"Python version: {python_version_single_line}")
    print(f"Weights file: {WEIGHTS_FILE}")
    print(f"Weights file exists: {WEIGHTS_FILE.exists()}")
    
    if WEIGHTS_FILE.exists():
        try:
            print(f"Weights file size: {WEIGHTS_FILE.stat().st_size} bytes")
        except Exception as e:
            print(f"Could not get weights file size: {e}")

    is_container = Path('/.dockerenv').exists() or 'DOCKER_HOST' in os.environ
    print(f"Running in container: {is_container}")

    if is_container:
        try:
            user_info = f"{os.getuid()}:{os.getgid()}"
            print(f"User running process: {user_info}")
        except AttributeError:
            print("User running process: UID/GID not available on this OS")
        
        for path_obj in [UPLOAD_FOLDER, RESULTS_FOLDER, ANNOT_FOLDER]:
            if path_obj.exists():
                stat_info = path_obj.stat()
                permissions = oct(stat_info.st_mode)[-3:]
                owner = f"{stat_info.st_uid}:{stat_info.st_gid}"
                print(f"Permissions for {path_obj}: {permissions}")
                print(f"Owner for {path_obj}: {owner}")
            else:
                print(f"Directory {path_obj} does not exist.")

    # some cleanup steps - not sure quite where to put these
    print('Running periodic cleanup of old sessions...')
    # Cleanup old session folders
    max_age_hours = 4
    now = time.time()
    for base_dir in [UPLOAD_FOLDER, RESULTS_FOLDER, ANNOT_FOLDER]:
        for session_dir in Path(base_dir).iterdir():
            if session_dir.is_dir():
                mtime = session_dir.stat().st_mtime
                if now - mtime > max_age_hours * 3600:
                    shutil.rmtree(session_dir)
    
    print('App is running at the following local addresses:',
          'http://127.0.0.1:7860⁠',
          'http://localhost:7860⁠',
          sep='\n')

if __name__ == '__main__':
    print_startup_info()
    app.run(host='0.0.0.0', port=7860, debug=True)