#!/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.5) 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)