WorkerPool threading and GPU support

app.py

@@ -6,33 +6,35 @@ import io
 import zipfile
 import cv2
 import csv
-import torch
+import pickle
 import json
 import shutil
+from ultralytics import YOLO
+from ultralytics.utils import ThreadingLocked
 import numpy as np
 import pandas as pd
+from torch import cuda
 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
+from yolo_utils import detect_in_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'
-RESULT_FOLDER = APP_ROOT / 'results'
+RESULTS_FOLDER = APP_ROOT / 'results'
 WEIGHTS_FILE = APP_ROOT / 'weights.pt'
 app.config['UPLOAD_FOLDER'] = str(UPLOAD_FOLDER)
-app.config['RESULT_FOLDER'] = str(RESULT_FOLDER)
+app.config['RESULTS_FOLDER'] = str(RESULTS_FOLDER)
 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)
+RESULTS_FOLDER.mkdir(parents=True, exist_ok=True)
 
 @app.errorhandler(Exception)
 def handle_exception(e):
@@ -47,13 +49,14 @@ def allowed_file(filename):
 def index():
     return render_template('index.html')
 
+
 # Load model once at startup, use CUDA if available
-MODEL_DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+MODEL_DEVICE = 'cuda' if cuda.is_available() else 'cpu'
 _model = None
 def get_model():
     global _model
     if _model is None:
-        _model = load_model(WEIGHTS_FILE)
+        _model = YOLO(WEIGHTS_FILE)
         if MODEL_DEVICE == 'cuda':
             _model.to('cuda')
     return _model
@@ -61,22 +64,23 @@ def get_model():
 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]:
+    results_dir = Path(app.config['RESULTS_FOLDER']) / session_id
+    for d in [upload_dir, results_dir]:
         if d.exists():
             shutil.rmtree(d)
 
 # save the uploaded files
-@app.route('/uploads/<session_id>', methods=['POST'])
-def upload_files(session_id):
+@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 = {}
-    file_data = []
     for f in files:
         orig_name = secure_filename(f.filename)
         ext = Path(orig_name).suffix
@@ -84,107 +88,164 @@ def upload_files(session_id):
         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
-
+    session['filename_map'] = filename_map
+    return jsonify({'filename_map': filename_map, 'status': 'uploaded'})
+
+# helper function to simplify args for pool.imap
+@ThreadingLocked()
+def process_single_image(args):
+    orig_name, img_path, pickle_path, model = args
+    img_results = detect_in_image(model, str(img_path))
+    with open(pickle_path, 'wb') as pf:
+        pickle.dump(img_results, pf)
+    return (orig_name, img_results)
 
-def process_image(args):
-    orig_name, unique_name, image_bytes = args
+@app.route('/process', methods=['POST'])
+def process_images():
     model = get_model()
-    detections = detect_image(model, image_bytes, conf=0.05)
-    return {'orig_name': orig_name, 'unique_name': unique_name, 'detections': detections}
-
-def async_process_images(session_id, file_data, state):
+    session_id = session['id']
+    filename_map = session.get('filename_map', {})
+    upload_dir = Path(app.config['UPLOAD_FOLDER']) / session_id
+    state = {}
+    state['status'] = 'starting'
+    state['progress'] = 0
+    state['filename_map'] = filename_map
+    state['jobId'] = session['id']
+    session['job_state'] = state
+
+    # create a results_dir, clean out old one if needed
+    results_dir = Path(app.config['RESULTS_FOLDER']) / session_id
+    if results_dir.exists():
+        shutil.rmtree(results_dir)
+    results_dir.mkdir(parents=True)
+
+    # set up args list for imap
+    n_img = len(filename_map)
+    arg_list = [(orig_name,
+                 upload_dir / filename_map[orig_name],
+                 results_dir / f"{Path(orig_name).stem}_results.pkl",
+                 model) for orig_name in filename_map.keys()]
     try:
-        state['status'] = 'running'
-        state['progress'] = 0
-        total = len(file_data)
-        results = []
-        detections = {}
-        # Use ThreadPool for CPU, else single-threaded for CUDA
+        all_detections = {}
+        state['status'] = 'processing'
+        session['job_state'] = state
         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']
-                state['progress'] = int((idx + 1) / total * 100)
+            for idx, args in enumerate(arg_list):
+                orig_name, img_results = process_single_image(args)
+                all_detections[orig_name] = img_results
+                state['progress'] = int((idx + 1) / n_img * 100)
+                session['job_state'] = state
         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
+                for idx, result in enumerate(pool.imap(process_single_image, arg_list)):
+                    state['progress'] = int((idx + 1) / n_img * 100)
+                    orig_name, img_results = result
+                    all_detections[orig_name] = img_results
+                    session['job_state'] = state
+        # Save all detections to a pickled file
+        detections_path = results_dir / 'all_detections.pkl'
+        with open(detections_path, 'wb') as f:
+            pickle.dump(all_detections, f)
+        state['status'] = 'completed'
         state['progress'] = 100
+        session['job_state'] = state
     except Exception as e:
+        print(f"Error in /process: {e}")
+        print(traceback.format_exc())
         state['status'] = 'error'
         state['error'] = str(e)
         state['progress'] = 100
-
-
-@app.route('/process', methods=['POST'])
-def process_images():
+        session['job_state'] = state
+    resp = {
+        'status': state.get('status', 'unknown'),
+        'progress': state.get('progress', 0),
+        'jobId': state.get('jobId'),
+        'error': state.get('error'),
+    }
+    return jsonify(resp)
+
+# Support /progress/<jobId> for frontend polling
+@app.route('/progress/<jobId>')
+def get_progress_with_id(jobId):
     try:
-        files = request.files.getlist('files')
-        if not files or files[0].filename == '':
-            return jsonify({'error': 'No files uploaded'}), 400
-        # 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': {},
+        job_state = session.get('job_state')
+        if not job_state:
+            print(f"/progress/{jobId}: No job_state found in session.")
+            return jsonify({"status": "error", "error": "No job state"}), 404
+        resp = {
+            'status': job_state.get('status', 'unknown'),
+            'progress': job_state.get('progress', 0),
+            'jobId': session.get('id'),
+            'error': job_state.get('error'),
         }
-        session['job_state'] = state
-        thread = Thread(target=async_process_images, args=(session_id, file_data, state))
-        thread.daemon = True
-        thread.start()
-        return jsonify({'jobId': session_id})
+        # If completed, load and return all_detections.pkl as JSON
+        if job_state.get('status') == 'completed':
+            session_id = session['id']
+            detections_path = Path(app.config['RESULTS_FOLDER']) / session_id / 'all_detections.pkl'
+            if detections_path.exists():
+                with open(detections_path, 'rb') as f:
+                    all_detections = pickle.load(f)
+                resp['results'] = all_detections
+        return jsonify(resp)
     except Exception as e:
-        print(f"Error in /process: {e}")
+        print(f"Error in /progress/{jobId}: {e}")
         print(traceback.format_exc())
-        return jsonify({'error': str(e)}), 500
+        return jsonify({"status": "error", "error": str(e)}), 500
 
-@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": "No job state"}), 404
-    if 'detections' in job_state:
-        job_state['detections_by_filename'] = job_state['detections']
-    return jsonify(job_state)
-
-@app.route('/results/<session_id>/<path:filename>')
-def download_file(session_id, filename):
+# /annotate route for dynamic annotation
+@app.route('/annotate', methods=['POST'])
+def annotate_image():
     try:
+        data = request.get_json()
+        filename = data.get('filename')
+        confidence = float(data.get('confidence', 0.5))
+        session_id = session['id']
+        filename_map = session.get('filename_map', {})
+        unique_name = filename_map.get(filename)
+        if not unique_name:
+            return jsonify({'error': 'File not found'}), 404
+        # Load detections from pickle
+        result_path = Path(app.config['RESULTS_FOLDER']) / session_id / f"{Path(filename).stem}_results.pkl"
+        if not result_path.exists():
+            return jsonify({'error': 'Results not found'}), 404
+        with open(result_path, 'rb') as pf:
+            detections = pickle.load(pf)
+        # Load image
+        img_path = Path(app.config['UPLOAD_FOLDER']) / session_id / unique_name
+        img = cv2.imread(str(img_path), cv2.IMREAD_UNCHANGED)
+        # Filter detections
+        filtered = [d for d in detections if d.get('score', 0) >= confidence]
+        # Draw boxes
+        for det in filtered:
+            x1, y1, x2, y2 = map(int, det['bbox'])
+            cv2.rectangle(img, (x1, y1), (x2, y2), (0,0,255), 3)
+        # Save annotated image to temp
+        annotated_path = Path(app.config['RESULTS_FOLDER']) / 'annotated'
+        annotated_path.mkdir(parents=True, exist_ok=True)
+        out_name = f"{Path(filename).stem}_annotated.png"
+        out_file = annotated_path / out_name
+        cv2.imwrite(str(out_file), img)
+        # Serve image
+        with open(out_file, 'rb') as f:
+            return send_file(
+                io.BytesIO(f.read()),
+                mimetype='image/png',
+                as_attachment=False,
+                download_name=out_name
+            )
+    except Exception as e:
+        print(f"Error in /annotate: {e}")
+        return jsonify({'error': str(e)}), 500
 
-        if 'id' not in session or session['id'] != session_id:
-            return jsonify({"error": "Session not found or expired"}), 404
+@app.route('/results/<path:filename>')
+def download_file(filename):
+    try:
+        session_id = session['id']
         if '..' in filename or filename.startswith('/'):
             return jsonify({"error": "Invalid filename"}), 400
         safe_filename = secure_filename(filename)
-        file_dir = Path(app.config['RESULT_FOLDER']) / session_id
+        file_dir = Path(app.config['RESULTS_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: {session_id}/{filename}")
@@ -235,13 +296,11 @@ def download_file(session_id, filename):
         print(error_message)
         return jsonify({"error": "Server error", "log": error_message}), 500
 
-@app.route('/export_images/<session_id>')
-def export_images(session_id):
+@app.route('/export_images')
+def export_images():
     try:
-
-        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
+        session_id = session['id']
+        job_dir = Path(app.config['RESULTS_FOLDER']) / session_id
         if not job_dir.exists():
             return jsonify({"error": f"Session directory {session_id} not found"}), 404
 
@@ -273,7 +332,7 @@ def export_images(session_id):
 def export_csv():
     try:
         data = request.json
-        session_id = session.get('id')
+        session_id = session['id']
         threshold = float(data.get('confidence', 0.5))
         job_state = session.get('job_state')
         if not job_state:
@@ -304,7 +363,7 @@ def export_csv():
 def export_images_post():
     try:
         data = request.json
-        session_id = session.get('id')
+        session_id = session['id']
         threshold = float(data.get('confidence', 0.5))
         job_state = session.get('job_state')
         if not job_state:
@@ -335,6 +394,8 @@ def export_images_post():
         print(error_message)
         return jsonify({"error": "Server error", "log": error_message}), 500
 
+
+
 def print_startup_info():
     print("----- NemaQuant Flask App Starting -----")
     print(f"Working directory: {os.getcwd()}")
@@ -384,11 +445,14 @@ 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:
+# Explicit endpoint for safe session cleanup
+@app.route('/cleanup', methods=['POST'])
+def cleanup_endpoint():
+    if 'id' in session:
         cleanup_session(session['id'])
+        session.clear()
+        return jsonify({'status': 'cleaned up'})
+    return jsonify({'error': 'No session to clean up'}), 400
 
 if __name__ == '__main__':
     print_startup_info()

static/script.js

@@ -35,6 +35,7 @@ document.addEventListener('DOMContentLoaded', () => {
     const MAX_ZOOM = 3;
     const MIN_ZOOM = 0.5;
     let progressInterval = null; // Interval timer for polling
+    let filteredValidFiles = [];
     
     // Panning variables
     let isPanning = false;
@@ -256,8 +257,34 @@ document.addEventListener('DOMContentLoaded', () => {
         fileInput.click();
     });
 
-    fileInput.addEventListener('change', () => {
+    fileInput.addEventListener('change', async () => {
         handleFiles(fileInput.files);
+        if (filteredValidFiles && filteredValidFiles.length > 0) {
+            // Prepare FormData for upload
+            const formData = new FormData();
+            filteredValidFiles.forEach(f => formData.append('files', f));
+            try {
+                const response = await fetch('/uploads', {
+                    method: 'POST',
+                    body: formData
+                });
+                if (response.ok) {
+                    const data = await response.json();
+                    logStatus('Files uploaded successfully.');
+                    // Update results table with filenames
+                    resultsTableBody.innerHTML = '';
+                    Object.keys(data.filename_map).forEach((filename, idx) => {
+                        const row = resultsTableBody.insertRow();
+                        row.dataset.originalIndex = idx;
+                        row.innerHTML = `<td>${filename}</td><td>Pending</td>`;
+                    });
+                } else {
+                    logStatus('File upload failed.');
+                }
+            } catch (err) {
+                logStatus('Error uploading files: ' + err);
+            }
+        }
     });
 
     // Input mode change
@@ -279,6 +306,24 @@ document.addEventListener('DOMContentLoaded', () => {
         } else {
             uploadText.textContent = `${validFileCount} image${validFileCount === 1 ? '' : 's'} selected`;
             startProcessingBtn.disabled = validFileCount === 0;
+            // Populate results table with filenames after upload
+            resultsTableBody.innerHTML = '';
+            filteredValidFiles.forEach((file, idx) => {
+                const row = resultsTableBody.insertRow();
+                row.dataset.originalIndex = idx;
+                row.innerHTML = `
+                    <td>${file.name}</td>
+                    <td>Pending</td>
+                    <td><button class="view-button" data-index="${idx}">View</button></td>
+                `;
+            });
+            // Add click event for View buttons
+            resultsTableBody.querySelectorAll('.view-button').forEach(btn => {
+                btn.addEventListener('click', (e) => {
+                    const idx = parseInt(btn.dataset.index, 10);
+                    displayImage(idx);
+                });
+            });
         }
     }
 
@@ -345,6 +390,7 @@ document.addEventListener('DOMContentLoaded', () => {
             }
             // --- ASYNC JOB: Start polling for progress ---
             if (data.jobId) {
+                logStatus(`Processing started. Job ID: ${data.jobId}`);
                 currentJobId = data.jobId;
                 pollProgress(currentJobId);
             } else {
@@ -365,27 +411,22 @@ document.addEventListener('DOMContentLoaded', () => {
     // --- Filtering and Table Update ---
     function updateResultsTable() {
         const threshold = parseFloat(confidenceSlider.value);
-        // Use allDetections (now an array of {filename, detections})
-        const grouped = {};
-        allDetections.forEach(imgResult => {
+        // Use allDetections (array of {filename, detections}) for filtering
+        const prevFilename = (currentImageIndex >= 0 && currentResults[currentImageIndex]) ? currentResults[currentImageIndex].filename : null;
+        currentResults = allDetections.map(imgResult => {
             const filtered = imgResult.detections.filter(det => det.score >= threshold);
-            grouped[imgResult.filename] = filtered;
+            return {
+                filename: imgResult.filename,
+                num_eggs: filtered.length,
+                detections: filtered
+            };
         });
-        // Build results for table
-        const prevFilename = (currentImageIndex >= 0 && currentResults[currentImageIndex]) ? currentResults[currentImageIndex].filename : null;
-        currentResults = Object.keys(grouped).map(filename => ({
-            filename,
-            num_eggs: grouped[filename].length,
-            detections: grouped[filename]
-        }));
         resultsTableBody.innerHTML = '';
         currentSortField = null;
         currentSortDirection = 'asc';
         totalPages = Math.ceil(currentResults.length / RESULTS_PER_PAGE);
         currentPage = 1;
         displayResultsPage(currentPage);
-        exportCsvBtn.disabled = true;
-        exportImagesBtn.disabled = true;
         // Try to restore previous image if it still exists
         let newIndex = 0;
         if (prevFilename) {
@@ -406,18 +447,39 @@ document.addEventListener('DOMContentLoaded', () => {
     }
 
     confidenceSlider.addEventListener('input', () => {
-        confidenceValue.textContent = confidenceSlider.value;
-        updateResultsTable();
-        renderConfidencePlot();
-        if (currentImageIndex >= 0) displayImage(currentImageIndex);
+    confidenceValue.textContent = confidenceSlider.value;
+    // 1. Update total eggs detected
+    const totalEggsElem = document.getElementById('total-eggs-count');
+    if (totalEggsElem && allDetections.length > 0) {
+        const threshold = parseFloat(confidenceSlider.value);
+        const totalEggs = allDetections.reduce((sum, imgResult) => sum + imgResult.detections.filter(det => det.score >= threshold).length, 0);
+        totalEggsElem.textContent = totalEggs;
+    }
+
+    // 2. Redraw confidence plot
+    renderConfidencePlot();
+
+    // 3. Update results table
+    updateResultsTable();
+
+    // 4. Update eggs detected under image preview
+    if (currentImageIndex >= 0) {
+        displayImage(currentImageIndex);
+    }
     });
 
     // --- Replace displayImage to use backend-annotated PNG ---
     async function displayImage(index) {
-        if (!currentResults[index]) return;
+        // Use either currentResults or filteredValidFiles for filename
+        let filename;
+        if (currentResults && currentResults[index] && currentResults[index].filename) {
+            filename = currentResults[index].filename;
+        } else if (filteredValidFiles && filteredValidFiles[index]) {
+            filename = filteredValidFiles[index].name;
+        } else {
+            return;
+        }
         currentImageIndex = index;
-        const result = currentResults[index];
-        const filename = result.filename;
         const confidence = parseFloat(confidenceSlider.value);
         try {
             const response = await fetch('/annotate', {
@@ -454,7 +516,7 @@ document.addEventListener('DOMContentLoaded', () => {
             clearPreview();
         }
         prevBtn.disabled = index <= 0;
-        nextBtn.disabled = index >= currentResults.length - 1;
+        nextBtn.disabled = index >= (currentResults.length > 0 ? currentResults.length : filteredValidFiles.length) - 1;
     }
 
     // --- New Polling Function ---
@@ -463,65 +525,65 @@ document.addEventListener('DOMContentLoaded', () => {
             clearInterval(progressInterval); // Clear any existing timer
         }
 
-        const totalImages = Array.from(fileInput.files).filter(file => {
-            const allowedTypes = ['image/png', 'image/jpeg', 'image/jpg', 'image/tiff', 'image/tif'];
-            if (inputMode.value === 'folder' || inputMode.value === 'keyence') {
-                return allowedTypes.includes(file.type) && 
-                       file.webkitRelativePath && 
-                       !file.webkitRelativePath.startsWith('.');
-            }
-            return allowedTypes.includes(file.type);
-        }).length;
-
         progressInterval = setInterval(async () => {
             try {
                 const response = await fetch(`/progress/${jobId}`);
                 if (!response.ok) {
-                     // Handle cases where the progress endpoint itself fails
-                     let errorText = `Progress check failed: ${response.status}`;
-                     try {
-                         const errorData = await response.json();
-                         errorText += `: ${errorData.error || 'Unknown progress error'}`;
-                     } catch(e) { errorText += ` - ${response.statusText}`; }
+                    let errorText = `Progress check failed: ${response.status}`;
+                    try {
+                        const errorData = await response.json();
+                        errorText += `: ${errorData.error || 'Unknown progress error'}`;
+                    } catch(e) { errorText += ` - ${response.statusText}`; }
                     throw new Error(errorText);
                 }
 
                 const data = await response.json();
-
-                // Update UI based on status
-                updateProgress(data.progress || 0, data.status.charAt(0).toUpperCase() + data.status.slice(1));
-                
-                // Update image counter based on progress percentage
-                if (data.status === 'Running' || data.status === 'Starting') {
-                    const processedImages = Math.floor((data.progress / 90) * totalImages); // 90 is the max progress before completion
-                    imageCounter.textContent = `${processedImages} of ${totalImages} images`;
-                }
-
-                if (data.status === 'success') {
-                    clearInterval(progressInterval);
-                    progressInterval = null;
-                    updateProgress(100, 'Processing complete');
-                    logStatus("Processing finished successfully.");
-                    displayResults(data.results || [], data.detections_by_filename || null);
-                    renderConfidencePlot();
-                    setLoading(false);
-                } else if (data.status === 'error') {
-                    clearInterval(progressInterval);
-                    progressInterval = null;
-                    logStatus(`Error during processing: ${data.error || 'Unknown error'}`);
-                    updateProgress(data.progress || 100, 'Error'); // Show error state on progress bar
-                    setLoading(false);
-                } else if (data.status === 'running' || data.status === 'starting') {
-                    // Continue polling
-                     // logStatus(`Processing status: ${data.status} (${data.progress}%)`); // Removed for cleaner log
-                     // You could add snippets from data.log here if the backend provides useful intermediate logs
-                } else {
-                     // Unknown status - stop polling to prevent infinite loops
-                     clearInterval(progressInterval);
-                     progressInterval = null;
-                     logStatus(`Warning: Unknown job status received: ${data.status}. Stopping progress updates.`);
-                     updateProgress(data.progress || 0, `Unknown (${data.status})`);
-                     setLoading(false); 
+                const status = (data.status || '').toLowerCase();
+
+                switch (status) {
+                    case 'starting':
+                        updateProgress(data.progress || 0, 'Starting...');
+                        logStatus('Job is starting.');
+                        break;
+                    case 'processing':
+                        updateProgress(data.progress || 0, `Processing (${data.progress || 0}%)`);
+                        logStatus(`Processing images... (${data.progress || 0}%)`);
+                        // If results are present, update detections and table
+                        if (data.results) {
+                            allDetections = Object.entries(data.results).map(([filename, detections]) => ({ filename, detections }));
+                            updateResultsTable();
+                        }
+                        break;
+                    case 'completed':
+                        clearInterval(progressInterval);
+                        progressInterval = null;
+                        updateProgress(100, 'Processing complete');
+                        logStatus('Processing finished successfully.');
+                        if (data.results) {
+                            allDetections = Object.entries(data.results).map(([filename, detections]) => ({ filename, detections }));
+                            updateResultsTable();
+                        }
+                        renderConfidencePlot();
+                        setLoading(false);
+                        break;
+                    case 'error':
+                        clearInterval(progressInterval);
+                        progressInterval = null;
+                        logStatus(`Error during processing: ${data.error || 'Unknown error'}`);
+                        updateProgress(data.progress || 100, 'Error');
+                        setLoading(false);
+                        break;
+                    case 'unknown':
+                        clearInterval(progressInterval);
+                        progressInterval = null;
+                        logStatus('Unknown job status. Stopping progress updates.');
+                        updateProgress(data.progress || 0, 'Unknown status');
+                        setLoading(false);
+                        break;
+                    default:
+                        logStatus(`Status: ${data.status}`);
+                        updateProgress(data.progress || 0, data.status ? data.status.charAt(0).toUpperCase() + data.status.slice(1) : '');
+                        break;
                 }
 
             } catch (error) {
@@ -531,7 +593,7 @@ document.addEventListener('DOMContentLoaded', () => {
                 updateProgress(0, 'Polling Error');
                 setLoading(false);
             }
-        }, 2000); // Poll every 2 seconds
+        }, 1000); // Poll every 1 second
     }
 
     // --- UI Update Functions ---
@@ -614,49 +676,40 @@ document.addEventListener('DOMContentLoaded', () => {
     }
 
     // Results Display
-    function displayResults(results, detectionsByFilename) {
-        // If detectionsByFilename is provided, update allDetections for plot logic
-        if (detectionsByFilename) {
-            allDetections = Object.keys(detectionsByFilename).map(filename => ({
-                filename,
-                detections: detectionsByFilename[filename]
-            }));
-        } else {
-            // Fallback: use results array
-            allDetections = results.map(r => ({ filename: r.filename, detections: r.detections }));
-        }
-        // Now update the table
-        currentResults = results;
+    function displayResults(jobStatus, filenameMap, resultsObj) {
         resultsTableBody.innerHTML = '';
         currentImageIndex = -1;
         currentSortField = null;
         currentSortDirection = 'asc';
 
-        if (!results || results.length === 0) {
-            logStatus("No results to display");
-            clearPreview();
+        // If job is not completed, show filenames only
+        if (jobStatus !== 'completed') {
+            Object.keys(filenameMap).forEach((filename, idx) => {
+                const row = resultsTableBody.insertRow();
+                row.innerHTML = `<td>${filename}</td><td>Pending</td>`;
+            });
+            exportCsvBtn.disabled = true;
+            exportImagesBtn.disabled = true;
+            logStatus('Waiting for job to complete...');
             return;
         }
 
-        // Reset sort indicators
-        document.querySelectorAll('.results-table th[data-sort]').forEach(h => {
-            h.classList.remove('sort-asc', 'sort-desc');
-        });
-
-        // Calculate pagination
-        totalPages = Math.ceil(results.length / RESULTS_PER_PAGE);
-        currentPage = 1;
-        
-        // Display current page
-        displayResultsPage(currentPage);
-
-        // Enable export buttons
-        exportCsvBtn.disabled = false;
-        exportImagesBtn.disabled = false;
-
-        // Show first image
-        displayImage(0);
-        logStatus(`Displayed ${results.length} results`);
+        // If job is completed, show filtered detection counts
+        if (resultsObj) {
+            Object.keys(resultsObj).forEach((filename, idx) => {
+                const detections = resultsObj[filename] || [];
+                // Filter by confidence threshold
+                const threshold = parseFloat(confidenceSlider.value);
+                const filtered = detections.filter(d => d.score >= threshold);
+                const row = resultsTableBody.insertRow();
+                row.innerHTML = `<td>${filename}</td><td>${filtered.length}</td>`;
+            });
+            exportCsvBtn.disabled = false;
+            exportImagesBtn.disabled = false;
+            logStatus('Job completed. Results displayed.');
+        } else {
+            logStatus('No results found.');
+        }
     }
     
     // Display a specific page of results
@@ -671,31 +724,33 @@ document.addEventListener('DOMContentLoaded', () => {
         for (let i = startIndex; i < endIndex; i++) {
             const result = currentResults[i];
             const row = resultsTableBody.insertRow();
-            row.innerHTML = `
-                <td>
-                    <i class="ri-image-line"></i>
-                    ${result.filename}
-                </td>
-                <td>${result.num_eggs}</td>
-                <td class="text-right">
-                    <button class="view-button" title="Click to view image">
-                        <i class="ri-eye-line"></i>
-                        View
-                    </button>
-                </td>
-            `;
-            
-            // Store the original index to maintain image preview relationship
-            row.dataset.originalIndex = currentResults.indexOf(result);
-            
-            row.addEventListener('click', () => {
-                const originalIndex = parseInt(row.dataset.originalIndex);
-                displayImage(originalIndex);
-                document.querySelectorAll('.results-table tr').forEach(r => r.classList.remove('selected'));
-                row.classList.add('selected');
+                row.innerHTML = `
+                    <td>
+                        <i class="ri-image-line"></i>
+                        ${result.filename}
+                    </td>
+                    <td>${result.num_eggs}</td>
+                    <td class="text-right">
+                        <button class="view-button" data-index="${i}" title="Click to view image">
+                            <i class="ri-eye-line"></i>
+                            View
+                        </button>
+                    </td>
+                `;
+                // Store the original index to maintain image preview relationship
+                row.dataset.originalIndex = i;
+            }
+        // Wire up View buttons after rows are created
+        resultsTableBody.querySelectorAll('.view-button').forEach(btn => {
+            btn.addEventListener('click', (e) => {
+                e.stopPropagation();
+                const idx = parseInt(btn.getAttribute('data-index'));
+                displayImage(idx);
+                // Highlight selected row
+                resultsTableBody.querySelectorAll('tr').forEach(r => r.classList.remove('selected'));
+                btn.closest('tr').classList.add('selected');
             });
-        }
-        
+        });
         // Update pagination UI
         updatePaginationControls();
     }

yolo_utils.py

@@ -1,14 +1,7 @@
-from ultralytics import YOLO
-import cv2
-import numpy as np
+from PIL import Image
 
-def load_model(weights_path):
-    return YOLO(weights_path)
-
-def detect_image(model, image_bytes, conf=0.05):
-    # image_bytes: bytes-like object (from Flask upload)
-    arr = np.frombuffer(image_bytes, np.uint8)
-    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
+def detect_in_image(model, im_path, conf=0.05):
+    img = Image.open(im_path)
     results = model.predict(img, imgsz=1440, max_det=1000, verbose=False, conf=conf)
     result = results[0]
     detections = []