recycling-project-solutions/main.py

#!/usr/bin/env python3
"""
Flask API for Memory Module Detection
This API processes motherboard images and detects memory modules using YOLOv8.
"""

import os
import io
import base64
from flask import Flask, request, jsonify, send_file, render_template
from flask_cors import CORS
from PIL import Image
import numpy as np
from werkzeug.utils import secure_filename
import tempfile
import logging
from inference_utils import MemoryModuleDetector

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Initialize Flask app
app = Flask(__name__)
CORS(app)

# Configuration
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max file size
UPLOAD_FOLDER = 'uploads'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'bmp'}

# Create upload folder if it doesn't exist
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# Initialize detector
MODEL_PATH = 'runs/detect/memory_module_detection/weights/best.pt'
detector = MemoryModuleDetector(MODEL_PATH)

# Hardcoded test image path
HARDCODED_IMAGE_PATH = 'training/memory/out1.png'

def allowed_file(filename):
    """Check if file extension is allowed."""
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def image_to_base64(image):
    """Convert PIL Image to base64 string."""
    buffer = io.BytesIO()
    image.save(buffer, format='PNG')
    img_str = base64.b64encode(buffer.getvalue()).decode()
    return img_str

def base64_to_image(base64_string):
    """Convert base64 string to PIL Image."""
    img_data = base64.b64decode(base64_string)
    image = Image.open(io.BytesIO(img_data))
    return image

@app.route('/', methods=['GET'])
def home():
    """Home endpoint - serve frontend or API information based on Accept header."""
    # Check if request is from a browser (wants HTML)
    if 'text/html' in request.headers.get('Accept', ''):
        return render_template('index.html')

    # Otherwise return JSON API information
    return jsonify({
        'message': 'Memory Module Detection API',
        'version': '1.0.0',
        'endpoints': {
            '/': 'GET - Frontend interface or API information',
            '/api': 'GET - API information (JSON)',
            '/detect': 'POST - Upload image for memory module detection',
            '/detect/hardcoded': 'GET - Process hardcoded test image',
            '/detect/base64': 'POST - Process base64 encoded image',
            '/health': 'GET - Health check'
        },
        'model_loaded': detector.model is not None,
        'supported_formats': list(ALLOWED_EXTENSIONS)
    })

@app.route('/api', methods=['GET'])
def api_info():
    """API information endpoint (always returns JSON)."""
    return jsonify({
        'message': 'Memory Module Detection API',
        'version': '1.0.0',
        'endpoints': {
            '/': 'GET - Frontend interface or API information',
            '/api': 'GET - API information (JSON)',
            '/detect': 'POST - Upload image for memory module detection',
            '/detect/hardcoded': 'GET - Process hardcoded test image',
            '/detect/base64': 'POST - Process base64 encoded image',
            '/health': 'GET - Health check'
        },
        'model_loaded': detector.model is not None,
        'supported_formats': list(ALLOWED_EXTENSIONS)
    })



@app.route('/health', methods=['GET'])
def health_check():
    """Health check endpoint."""
    return jsonify({
        'status': 'healthy',
        'model_loaded': detector.model is not None,
        'model_path': MODEL_PATH
    })

@app.route('/detect', methods=['POST'])
def detect_memory_modules():
    """
    Detect memory modules in uploaded image.
    
    Expected input:
    - File upload with key 'image'
    - Optional: confidence threshold as form data
    
    Returns:
    - JSON with detections and annotated image (base64)
    """
    try:
        # Check if model is loaded
        if detector.model is None:
            return jsonify({
                'error': 'Model not loaded. Please train the model first.',
                'success': False
            }), 500
        
        # Check if file is present
        if 'image' not in request.files:
            return jsonify({
                'error': 'No image file provided',
                'success': False
            }), 400
        
        file = request.files['image']
        
        # Check if file is selected
        if file.filename == '':
            return jsonify({
                'error': 'No file selected',
                'success': False
            }), 400
        
        # Check file extension
        if not allowed_file(file.filename):
            return jsonify({
                'error': f'File type not allowed. Supported formats: {ALLOWED_EXTENSIONS}',
                'success': False
            }), 400
        
        # Get confidence threshold from form data (default 80%)
        conf_threshold = float(request.form.get('confidence', 0.8))
        
        # Save uploaded file temporarily
        filename = secure_filename(file.filename)
        temp_path = os.path.join(UPLOAD_FOLDER, filename)
        file.save(temp_path)
        
        try:
            # Run detection
            detections, annotated_image = detector.detect(
                temp_path, 
                conf_threshold=conf_threshold
            )
            
            # Convert annotated image to base64
            annotated_base64 = image_to_base64(annotated_image)
            
            # Prepare response
            response_data = {
                'success': True,
                'detections': detections,
                'num_detections': len(detections),
                'annotated_image': annotated_base64,
                'confidence_threshold': conf_threshold,
                'original_filename': filename
            }
            
            logger.info(f"Processed {filename}: found {len(detections)} memory modules")
            
            return jsonify(response_data)
        
        finally:
            # Clean up temporary file
            if os.path.exists(temp_path):
                os.remove(temp_path)
    
    except Exception as e:
        logger.error(f"Error processing image: {str(e)}")
        return jsonify({
            'error': f'Error processing image: {str(e)}',
            'success': False
        }), 500

@app.route('/detect/hardcoded', methods=['GET'])
def detect_hardcoded_image():
    """
    Process hardcoded test image for memory module detection.

    Optional query parameters:
    - confidence: confidence threshold (default: 0.8)

    Returns:
    - JSON with detections and annotated image (base64)
    """
    try:
        # Check if model is loaded
        if detector.model is None:
            return jsonify({
                'error': 'Model not loaded. Please train the model first.',
                'success': False
            }), 500

        # Check if hardcoded image exists
        if not os.path.exists(HARDCODED_IMAGE_PATH):
            return jsonify({
                'error': f'Hardcoded test image not found at {HARDCODED_IMAGE_PATH}',
                'success': False
            }), 404

        # Get confidence threshold from query parameters (default 80%)
        conf_threshold = float(request.args.get('confidence', 0.8))

        # Run detection
        detections, annotated_image = detector.detect(
            HARDCODED_IMAGE_PATH,
            conf_threshold=conf_threshold
        )

        # Convert annotated image to base64
        annotated_base64 = image_to_base64(annotated_image)

        # Prepare response
        response_data = {
            'success': True,
            'detections': detections,
            'num_detections': len(detections),
            'annotated_image': annotated_base64,
            'confidence_threshold': conf_threshold,
            'test_image_path': HARDCODED_IMAGE_PATH
        }

        logger.info(f"Processed hardcoded image: found {len(detections)} memory modules")

        return jsonify(response_data)
    
    except Exception as e:
        logger.error(f"Error processing hardcoded image: {str(e)}")
        return jsonify({
            'error': f'Error processing hardcoded image: {str(e)}',
            'success': False
        }), 500



@app.route('/detect/base64', methods=['POST'])
def detect_base64_image():
    """
    Detect memory modules in base64 encoded image.
    
    Expected JSON input:
    {
        "image": "base64_encoded_image_string",
        "confidence": 0.5  // optional
    }
    
    Returns:
    - JSON with detections and annotated image (base64)
    """
    try:
        # Check if model is loaded
        if detector.model is None:
            return jsonify({
                'error': 'Model not loaded. Please train the model first.',
                'success': False
            }), 500
        
        # Get JSON data
        data = request.get_json()
        
        if not data or 'image' not in data:
            return jsonify({
                'error': 'No base64 image data provided',
                'success': False
            }), 400
        
        # Get confidence threshold (default 80%)
        conf_threshold = float(data.get('confidence', 0.8))
        
        # Decode base64 image
        try:
            image = base64_to_image(data['image'])
        except Exception as e:
            return jsonify({
                'error': f'Invalid base64 image data: {str(e)}',
                'success': False
            }), 400
        
        # Run detection
        detections, annotated_image = detector.detect_from_array(
            np.array(image), 
            conf_threshold=conf_threshold
        )
        
        # Convert annotated image to base64
        annotated_base64 = image_to_base64(annotated_image)
        
        # Prepare response
        response_data = {
            'success': True,
            'detections': detections,
            'num_detections': len(detections),
            'annotated_image': annotated_base64,
            'confidence_threshold': conf_threshold
        }
        
        logger.info(f"Processed base64 image: found {len(detections)} memory modules")
        
        return jsonify(response_data)
    
    except Exception as e:
        logger.error(f"Error processing base64 image: {str(e)}")
        return jsonify({
            'error': f'Error processing base64 image: {str(e)}',
            'success': False
        }), 500

@app.errorhandler(413)
def too_large(e):
    """Handle file too large error."""
    return jsonify({
        'error': 'File too large. Maximum size is 16MB.',
        'success': False
    }), 413

@app.errorhandler(404)
def not_found(e):
    """Handle 404 errors."""
    return jsonify({
        'error': 'Endpoint not found',
        'success': False
    }), 404

@app.errorhandler(500)
def internal_error(e):
    """Handle internal server errors."""
    return jsonify({
        'error': 'Internal server error',
        'success': False
    }), 500

if __name__ == '__main__':
    # Check if model exists
    if not os.path.exists(MODEL_PATH):
        print(f"Warning: Model not found at {MODEL_PATH}")
        print("Please train the model first using: python3 train.py")
        print("The API will still start but detection endpoints will return errors.")
    
    # Start the Flask app
    print("Starting Memory Module Detection API...")
    print(f"Model path: {MODEL_PATH}")
    print(f"Model loaded: {detector.model is not None}")
    print(f"Hardcoded test image: {HARDCODED_IMAGE_PATH}")
    
    app.run(host='0.0.0.0', port=5002, debug=True)