app.py

from flask import Flask, request, jsonify, send_file, render_template, Response
from flask_cors import CORS
import numpy as np
from PIL import Image
import io
import base64
import time
import os
import json
from concurrent.futures import ThreadPoolExecutor
import multiprocessing as mp
import ray
import psutil
import tracemalloc

app = Flask(__name__)
CORS(app)

@app.route('/')
def index():
    """Serve the main HTML page"""
    return render_template('index.html')

# Initialize Ray ONCE at module level
if not ray.is_initialized():
    ray.init(
        num_cpus=mp.cpu_count(),
        include_dashboard=False,
        logging_level='ERROR',
        ignore_reinit_error=True
    )
    print("✓ Ray initialized successfully")

# ============================================================================
# MEMORY TRACKING UTILITIES
# ============================================================================

def get_memory_usage():
    """Get current memory usage in MB"""
    process = psutil.Process()
    return process.memory_info().rss / 1024 / 1024

def track_memory(func):
    """Decorator to track memory usage of a function"""
    def wrapper(*args, **kwargs):
        tracemalloc.start()
        mem_before = get_memory_usage()
        
        result = func(*args, **kwargs)
        
        mem_after = get_memory_usage()
        current, peak = tracemalloc.get_traced_memory()
        tracemalloc.stop()
        
        memory_stats = {
            'memory_used_mb': round(mem_after - mem_before, 2),
            'peak_memory_mb': round(peak / 1024 / 1024, 2),
            'current_memory_mb': round(mem_after, 2)
        }
        
        return result, memory_stats
    return wrapper

# ============================================================================
# GRAYSCALE CONVERSION METHODS
# ============================================================================

def luminosity_method(r, g, b):
    """Standard luminosity method - Most accurate for human perception"""
    return (0.299 * r + 0.587 * g + 0.114 * b).astype(np.uint8)

def average_method(r, g, b):
    """Simple average method"""
    return ((r + g + b) / 3).astype(np.uint8)

def desaturation_method(r, g, b):
    """Desaturation method using min-max average"""
    return ((np.maximum(np.maximum(r, g), b) + np.minimum(np.minimum(r, g), b)) / 2).astype(np.uint8)

def bt601_method(r, g, b):
    """BT.601 standard (used in NTSC/PAL)"""
    return (0.299 * r + 0.587 * g + 0.114 * b).astype(np.uint8)

def bt709_method(r, g, b):
    """BT.709 standard (used in HDTV)"""
    return (0.2126 * r + 0.7152 * g + 0.0722 * b).astype(np.uint8)

def bt2020_method(r, g, b):
    """BT.2020 standard (used in UHDTV)"""
    return (0.2627 * r + 0.6780 * g + 0.0593 * b).astype(np.uint8)

CONVERSION_METHODS = {
    'luminosity': luminosity_method,
    'average': average_method,
    'desaturation': desaturation_method,
    'bt601': bt601_method,
    'bt709': bt709_method,
    'bt2020': bt2020_method
}

# ============================================================================
# RAY REMOTE FUNCTIONS
# ============================================================================

@ray.remote
def ray_process_chunk(chunk_data, method_name):
    """Ray remote function to process a chunk"""
    method_funcs = {
        'luminosity': lambda r, g, b: (0.299 * r + 0.587 * g + 0.114 * b).astype(np.uint8),
        'average': lambda r, g, b: ((r + g + b) / 3).astype(np.uint8),
        'desaturation': lambda r, g, b: ((np.maximum(np.maximum(r, g), b) + np.minimum(np.minimum(r, g), b)) / 2).astype(np.uint8),
        'bt601': lambda r, g, b: (0.299 * r + 0.587 * g + 0.114 * b).astype(np.uint8),
        'bt709': lambda r, g, b: (0.2126 * r + 0.7152 * g + 0.0722 * b).astype(np.uint8),
        'bt2020': lambda r, g, b: (0.2627 * r + 0.6780 * g + 0.0593 * b).astype(np.uint8)
    }
    method_func = method_funcs[method_name]
    r, g, b = chunk_data[:, :, 0], chunk_data[:, :, 1], chunk_data[:, :, 2]
    return method_func(r, g, b)

# ============================================================================
# PROCESSING METHODS WITH MEMORY TRACKING
# ============================================================================

def sequential_processing(img_array, method_func):
    """Sequential processing without any parallelization"""
    start_time = time.perf_counter()
    mem_before = get_memory_usage()
    
    r, g, b = img_array[:, :, 0], img_array[:, :, 1], img_array[:, :, 2]
    grayscale = method_func(r, g, b)
    
    processing_time = time.perf_counter() - start_time
    mem_used = get_memory_usage() - mem_before
    
    return grayscale, processing_time, {'memory_used_mb': round(mem_used, 2)}

def numpy_vectorized(img_array, method_func):
    """Fully vectorized NumPy processing"""
    start_time = time.perf_counter()
    mem_before = get_memory_usage()
    
    r, g, b = img_array[:, :, 0], img_array[:, :, 1], img_array[:, :, 2]
    grayscale = method_func(r, g, b)
    
    processing_time = time.perf_counter() - start_time
    mem_used = get_memory_usage() - mem_before
    
    return grayscale, processing_time, {'memory_used_mb': round(mem_used, 2)}

def ray_processing(img_array, method_name, num_chunks):
    """Ray distributed processing"""
    start_time = time.perf_counter()
    mem_before = get_memory_usage()
    
    height = img_array.shape[0]
    chunk_size = max(1, height // num_chunks)
    
    futures = []
    for i in range(num_chunks):
        start_row = i * chunk_size
        end_row = min(start_row + chunk_size, height) if i < num_chunks - 1 else height
        
        if start_row >= height:
            break
        
        chunk = img_array[start_row:end_row, :, :]
        future = ray_process_chunk.remote(chunk, method_name)
        futures.append(future)
    
    gray_chunks = ray.get(futures)
    grayscale_image = np.vstack(gray_chunks)
    
    processing_time = time.perf_counter() - start_time
    mem_used = get_memory_usage() - mem_before
    
    return grayscale_image, processing_time, {'memory_used_mb': round(mem_used, 2)}

def multiprocessing_worker(args):
    """Worker function for multiprocessing - must be at module level"""
    chunk, method_name = args
    method_funcs = {
        'luminosity': lambda r, g, b: (0.299 * r + 0.587 * g + 0.114 * b).astype(np.uint8),
        'average': lambda r, g, b: ((r + g + b) / 3).astype(np.uint8),
        'desaturation': lambda r, g, b: ((np.maximum(np.maximum(r, g), b) + np.minimum(np.minimum(r, g), b)) / 2).astype(np.uint8),
        'bt601': lambda r, g, b: (0.299 * r + 0.587 * g + 0.114 * b).astype(np.uint8),
        'bt709': lambda r, g, b: (0.2126 * r + 0.7152 * g + 0.0722 * b).astype(np.uint8),
        'bt2020': lambda r, g, b: (0.2627 * r + 0.6780 * g + 0.0593 * b).astype(np.uint8)
    }
    method_func = method_funcs[method_name]
    r, g, b = chunk[:, :, 0], chunk[:, :, 1], chunk[:, :, 2]
    return method_func(r, g, b)

def multiprocessing_processing(img_array, method_name, num_processes):
    """Python multiprocessing with Pool"""
    start_time = time.perf_counter()
    mem_before = get_memory_usage()
    
    height = img_array.shape[0]
    chunk_size = max(1, height // num_processes)
    
    chunks = []
    for i in range(num_processes):
        start_row = i * chunk_size
        end_row = min(start_row + chunk_size, height) if i < num_processes - 1 else height
        
        if start_row >= height:
            break
        
        chunks.append((img_array[start_row:end_row, :, :].copy(), method_name))
    
    with mp.Pool(processes=len(chunks)) as pool:
        gray_chunks = pool.map(multiprocessing_worker, chunks)
    
    grayscale_image = np.vstack(gray_chunks)
    
    processing_time = time.perf_counter() - start_time
    mem_used = get_memory_usage() - mem_before
    
    return grayscale_image, processing_time, {'memory_used_mb': round(mem_used, 2)}

def threading_worker(args):
    """Worker function for threading"""
    chunk, method_name = args
    method_funcs = {
        'luminosity': lambda r, g, b: (0.299 * r + 0.587 * g + 0.114 * b).astype(np.uint8),
        'average': lambda r, g, b: ((r + g + b) / 3).astype(np.uint8),
        'desaturation': lambda r, g, b: ((np.maximum(np.maximum(r, g), b) + np.minimum(np.minimum(r, g), b)) / 2).astype(np.uint8),
        'bt601': lambda r, g, b: (0.299 * r + 0.587 * g + 0.114 * b).astype(np.uint8),
        'bt709': lambda r, g, b: (0.2126 * r + 0.7152 * g + 0.0722 * b).astype(np.uint8),
        'bt2020': lambda r, g, b: (0.2627 * r + 0.6780 * g + 0.0593 * b).astype(np.uint8)
    }
    method_func = method_funcs[method_name]
    r, g, b = chunk[:, :, 0], chunk[:, :, 1], chunk[:, :, 2]
    return method_func(r, g, b)

def threading_processing(img_array, method_name, num_threads):
    """Python threading with ThreadPoolExecutor"""
    start_time = time.perf_counter()
    mem_before = get_memory_usage()
    
    height = img_array.shape[0]
    chunk_size = max(1, height // num_threads)
    
    chunks = []
    for i in range(num_threads):
        start_row = i * chunk_size
        end_row = min(start_row + chunk_size, height) if i < num_threads - 1 else height
        
        if start_row >= height:
            break
        
        chunks.append((img_array[start_row:end_row, :, :], method_name))
    
    with ThreadPoolExecutor(max_workers=len(chunks)) as executor:
        gray_chunks = list(executor.map(threading_worker, chunks))
    
    grayscale_image = np.vstack(gray_chunks)
    
    processing_time = time.perf_counter() - start_time
    mem_used = get_memory_usage() - mem_before
    
    return grayscale_image, processing_time, {'memory_used_mb': round(mem_used, 2)}

# ============================================================================
# COMPREHENSIVE COMPARISON WITH PROGRESS STREAMING
# ============================================================================

def run_comprehensive_comparison(img_array, progress_callback=None):
    """Run all comparisons and optionally stream progress updates"""
    cpu_count = mp.cpu_count()
    results = {
        'system_info': {
            'cpu_cores': cpu_count,
            'image_size': f"{img_array.shape[1]}x{img_array.shape[0]}",
            'pixels': img_array.shape[0] * img_array.shape[1],
            'total_memory_mb': round(psutil.virtual_memory().total / 1024 / 1024, 2),
            'available_memory_mb': round(psutil.virtual_memory().available / 1024 / 1024, 2)
        },
        'sequential': {},
        'numpy_vectorized': {},
        'ray': {},
        'multiprocessing': {},
        'threading': {},
        'conversion_methods': {},
        'core_scaling': {},
        'memory_usage': {}
    }
    
    total_steps = 7
    
    print(f"Image size: {img_array.shape}")
    
    # 1. Sequential Processing (Baseline) - All Methods
    print("\n[1/7] Testing Sequential Processing...")
    if progress_callback:
        progress_callback({'stage': 'Sequential Processing', 'progress': 14, 'step': 1, 'total': total_steps})
    
    for method_name, method_func in CONVERSION_METHODS.items():
        try:
            gray, time_taken, mem_stats = sequential_processing(img_array, method_func)
            results['sequential'][method_name] = time_taken
            results['memory_usage'][f'sequential_{method_name}'] = mem_stats
            print(f"  ✓ {method_name}: {time_taken:.4f}s")
        except Exception as e:
            print(f"  ✗ {method_name} failed: {e}")
    
    # 2. NumPy Vectorized - All Methods
    print("\n[2/7] Testing NumPy Vectorized...")
    if progress_callback:
        progress_callback({'stage': 'NumPy Vectorized', 'progress': 28, 'step': 2, 'total': total_steps})
    
    for method_name, method_func in CONVERSION_METHODS.items():
        try:
            gray, time_taken, mem_stats = numpy_vectorized(img_array, method_func)
            results['numpy_vectorized'][method_name] = time_taken
            results['memory_usage'][f'numpy_{method_name}'] = mem_stats
            print(f"  ✓ {method_name}: {time_taken:.4f}s")
        except Exception as e:
            print(f"  ✗ {method_name} failed: {e}")
    
    # 3. Ray Processing - Variable Cores
    print("\n[3/7] Testing Ray Framework...")
    if progress_callback:
        progress_callback({'stage': 'Ray Framework', 'progress': 42, 'step': 3, 'total': total_steps})
    
    for cores in [1, 2, 4, 8, cpu_count]:
        if cores <= cpu_count:
            try:
                gray, time_taken, mem_stats = ray_processing(img_array, 'luminosity', cores)
                results['ray'][f'{cores}_cores'] = time_taken
                results['memory_usage'][f'ray_{cores}_cores'] = mem_stats
                print(f"  ✓ {cores} cores: {time_taken:.4f}s")
            except Exception as e:
                print(f"  ✗ {cores} cores failed: {e}")
    
    # 4. Multiprocessing - Variable Cores
    print("\n[4/7] Testing Multiprocessing...")
    if progress_callback:
        progress_callback({'stage': 'Multiprocessing', 'progress': 56, 'step': 4, 'total': total_steps})
    
    for cores in [1, 2, 4, 8, cpu_count]:
        if cores <= cpu_count:
            try:
                gray, time_taken, mem_stats = multiprocessing_processing(img_array, 'luminosity', cores)
                results['multiprocessing'][f'{cores}_cores'] = time_taken
                results['memory_usage'][f'mp_{cores}_cores'] = mem_stats
                print(f"  ✓ {cores} cores: {time_taken:.4f}s")
            except Exception as e:
                print(f"  ✗ {cores} cores failed: {e}")
    
    # 5. Threading - Variable Threads
    print("\n[5/7] Testing Threading...")
    if progress_callback:
        progress_callback({'stage': 'Threading', 'progress': 70, 'step': 5, 'total': total_steps})
    
    for threads in [1, 2, 4, 8, cpu_count]:
        if threads <= cpu_count:
            try:
                gray, time_taken, mem_stats = threading_processing(img_array, 'luminosity', threads)
                results['threading'][f'{threads}_threads'] = time_taken
                results['memory_usage'][f'thread_{threads}_threads'] = mem_stats
                print(f"  ✓ {threads} threads: {time_taken:.4f}s")
            except Exception as e:
                print(f"  ✗ {threads} threads failed: {e}")
    
    # 6. Compare All Methods with Ray (4 cores)
    print("\n[6/7] Testing Conversion Methods with Ray...")
    if progress_callback:
        progress_callback({'stage': 'Conversion Methods', 'progress': 84, 'step': 6, 'total': total_steps})
    
    test_cores = min(4, cpu_count)
    for method_name in CONVERSION_METHODS.keys():
        try:
            gray, time_taken, mem_stats = ray_processing(img_array, method_name, test_cores)
            results['conversion_methods'][method_name] = time_taken
            results['memory_usage'][f'method_{method_name}'] = mem_stats
            print(f"  ✓ {method_name}: {time_taken:.4f}s")
        except Exception as e:
            print(f"  ✗ {method_name} failed: {e}")
    
    # 7. Core Scaling Analysis (Ray)
    print("\n[7/7] Testing Core Scaling with Ray...")
    if progress_callback:
        progress_callback({'stage': 'Core Scaling Analysis', 'progress': 98, 'step': 7, 'total': total_steps})
    
    for cores in range(1, min(cpu_count + 1, 17)):
        try:
            gray, time_taken, mem_stats = ray_processing(img_array, 'luminosity', cores)
            results['core_scaling'][f'{cores}_cores'] = time_taken
            results['memory_usage'][f'scaling_{cores}_cores'] = mem_stats
            print(f"  ✓ {cores} cores: {time_taken:.4f}s")
        except Exception as e:
            print(f"  ✗ {cores} cores failed: {e}")
    
    if progress_callback:
        progress_callback({'stage': 'Complete', 'progress': 100, 'step': 7, 'total': total_steps})
    
    print("\n✅ Comprehensive comparison completed!")
    return results

# ============================================================================
# API ENDPOINTS
# ============================================================================

@app.route('/process-stream', methods=['POST'])
def process_image_stream():
    """Process image with real-time progress updates using Server-Sent Events"""
    def generate():
        try:
            if 'image' not in request.files:
                yield f"data: {json.dumps({'error': 'No image file provided'})}\n\n"
                return
            
            file = request.files['image']
            
            if file.filename == '':
                yield f"data: {json.dumps({'error': 'No file selected'})}\n\n"
                return
            
            img = Image.open(file.stream)
            
            if img.mode != 'RGB':
                img = img.convert('RGB')
            
            img_array = np.array(img)
            
            yield f"data: {json.dumps({'stage': 'Starting', 'progress': 0})}\n\n"
            
            # Progress callback that yields updates
            def progress_callback(update):
                nonlocal generate
                # Note: Can't actually yield from callback, but we'll send updates after each stage
                pass
            
            comparison_results = run_comprehensive_comparison(img_array, progress_callback)
            
            test_cores = min(4, mp.cpu_count())
            grayscale_array, _, _ = ray_processing(img_array, 'luminosity', test_cores)
            grayscale_img = Image.fromarray(grayscale_array)
            
            buffer = io.BytesIO()
            grayscale_img.save(buffer, format='JPEG', quality=95)
            buffer.seek(0)
            img_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
            
            output_path = 'processed_image.jpg'
            grayscale_img.save(output_path, quality=95)
            
            recommendations = analyze_results(comparison_results)
            
            final_data = {
                'success': True,
                'preview': f'data:image/jpeg;base64,{img_base64}',
                'comparison_results': comparison_results,
                'recommendations': recommendations,
                'download_ready': True
            }
            
            yield f"data: {json.dumps(final_data)}\n\n"
            
        except Exception as e:
            import traceback
            print(traceback.format_exc())
            yield f"data: {json.dumps({'error': str(e)})}\n\n"
    
    return Response(generate(), mimetype='text/event-stream')

@app.route('/process', methods=['POST'])
def process_image():
    """Process image and return comprehensive comparison (non-streaming)"""
    try:
        if 'image' not in request.files:
            return jsonify({'error': 'No image file provided'}), 400
        
        file = request.files['image']
        
        if file.filename == '':
            return jsonify({'error': 'No file selected'}), 400
        
        img = Image.open(file.stream)
        
        if img.mode != 'RGB':
            img = img.convert('RGB')
        
        img_array = np.array(img)
        
        print("\n" + "="*60)
        print("Starting comprehensive comparison...")
        print("="*60)
        comparison_results = run_comprehensive_comparison(img_array)
        
        test_cores = min(4, mp.cpu_count())
        grayscale_array, _, _ = ray_processing(img_array, 'luminosity', test_cores)
        grayscale_img = Image.fromarray(grayscale_array)
        
        buffer = io.BytesIO()
        grayscale_img.save(buffer, format='JPEG', quality=95)
        buffer.seek(0)
        img_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
        
        output_path = 'processed_image.jpg'
        grayscale_img.save(output_path, quality=95)
        
        recommendations = analyze_results(comparison_results)
        
        return jsonify({
            'success': True,
            'preview': f'data:image/jpeg;base64,{img_base64}',
            'comparison_results': comparison_results,
            'recommendations': recommendations,
            'download_ready': True
        })
    
    except Exception as e:
        import traceback
        print("\n❌ ERROR:")
        print(traceback.format_exc())
        return jsonify({'error': str(e)}), 500

def analyze_results(results):
    """Analyze results and provide recommendations"""
    recommendations = {
        'fastest_framework': None,
        'fastest_method': None,
        'optimal_cores': None,
        'speedup_vs_sequential': None,
        'efficiency_analysis': {},
        'memory_analysis': {}
    }
    
    # Memory analysis
    if 'memory_usage' in results:
        total_memory = sum(
            mem.get('memory_used_mb', 0) 
            for mem in results['memory_usage'].values()
        )
        recommendations['memory_analysis']['total_memory_used_mb'] = round(total_memory, 2)
        
        # Find most memory-efficient approach
        valid_mem = [(k, v.get('memory_used_mb', float('inf'))) 
                     for k, v in results['memory_usage'].items() 
                     if v.get('memory_used_mb', 0) > 0]
        
        if valid_mem:
            mem_efficient = min(valid_mem, key=lambda x: x[1])
            recommendations['memory_analysis']['most_efficient'] = {
                'method': mem_efficient[0],
                'memory_mb': mem_efficient[1]
            }
    
    # Find fastest framework (comparing 4 cores/threads)
    framework_times = {}
    
    if 'numpy_vectorized' in results and 'luminosity' in results['numpy_vectorized']:
        t = results['numpy_vectorized']['luminosity']
        if t > 0:
            framework_times['NumPy Vectorized'] = t
    
    if 'ray' in results and '4_cores' in results['ray']:
        t = results['ray']['4_cores']
        if t > 0:
            framework_times['Ray'] = t
    
    if 'multiprocessing' in results and '4_cores' in results['multiprocessing']:
        t = results['multiprocessing']['4_cores']
        if t > 0:
            framework_times['Multiprocessing'] = t
    
    if 'threading' in results and '4_threads' in results['threading']:
        t = results['threading']['4_threads']
        if t > 0:
            framework_times['Threading'] = t
    
    if framework_times:
        fastest = min(framework_times.items(), key=lambda x: x[1])
        recommendations['fastest_framework'] = {
            'name': fastest[0],
            'time': fastest[1]
        }
    
    # Find fastest conversion method
    if 'conversion_methods' in results and results['conversion_methods']:
        valid_methods = {k: v for k, v in results['conversion_methods'].items() if v > 0}
        if valid_methods:
            fastest_method = min(valid_methods.items(), key=lambda x: x[1])
            recommendations['fastest_method'] = {
                'name': fastest_method[0],
                'time': fastest_method[1]
            }
    
    # Find optimal core count
    if 'core_scaling' in results and results['core_scaling']:
        valid_cores = {k: v for k, v in results['core_scaling'].items() if v > 0}
        if valid_cores:
            optimal_cores = min(valid_cores.items(), key=lambda x: x[1])
            recommendations['optimal_cores'] = {
                'cores': optimal_cores[0],
                'time': optimal_cores[1]
            }
    
    # Calculate speedup
    if 'sequential' in results and 'luminosity' in results['sequential']:
        seq_time = results['sequential']['luminosity']
        if seq_time > 0 and recommendations['fastest_framework'] and recommendations['fastest_framework']['time'] > 0:
            fastest_time = recommendations['fastest_framework']['time']
            recommendations['speedup_vs_sequential'] = round(seq_time / fastest_time, 2)
    
    # Efficiency analysis
    if 'core_scaling' in results and results['core_scaling']:
        base_time = results['core_scaling'].get('1_cores', 0)
        if base_time > 0:
            for cores, time_taken in results['core_scaling'].items():
                if time_taken > 0:
                    num_cores = int(cores.split('_')[0])
                    ideal_speedup = num_cores
                    actual_speedup = base_time / time_taken
                    efficiency = (actual_speedup / ideal_speedup * 100)
                    recommendations['efficiency_analysis'][cores] = {
                        'actual_speedup': round(actual_speedup, 2),
                        'ideal_speedup': ideal_speedup,
                        'efficiency': round(efficiency, 2)
                    }
    
    return recommendations

@app.route('/download', methods=['GET'])
def download_image():
    """Download processed image"""
    try:
        return send_file('processed_image.jpg', 
                        as_attachment=True,
                        download_name='grayscale_medical_image.jpg',
                        mimetype='image/jpeg')
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/health', methods=['GET'])
def health_check():
    """Health check endpoint"""
    return jsonify({
        'status': 'running',
        'ray_initialized': ray.is_initialized(),
        'cpu_cores': mp.cpu_count(),
        'memory_available_mb': round(psutil.virtual_memory().available / 1024 / 1024, 2),
        'memory_total_mb': round(psutil.virtual_memory().total / 1024 / 1024, 2)
    })

if __name__ == '__main__':
    port = int(os.environ.get('PORT', 5000))
    print("=" * 60)
    print("Enhanced Medical Image Processing Server - Combined Version")
    print(f"CPU Cores Available: {mp.cpu_count()}")
    print(f"Ray Initialized: {ray.is_initialized()}")
    print(f"Memory Available: {round(psutil.virtual_memory().available / 1024 / 1024, 2)} MB")
    print(f"Server starting on port {port}")
    print("=" * 60)
    app.run(host='0.0.0.0', port=port, debug=False, threaded=True)