Aherobo Ovie Victor
780e32c412
✅ Simplified Test Logic: - Removed unnecessary /detect/no-memory endpoint - Reverted to original 3 tests structure - Test 1: API Health Check - Test 2: Image with Memory Modules - Test 3: API Information ✅ Smart Message Display: - When memory modules found: '✅ Found X memory modules' - When no memory modules found: '❌ No memory modules' - Same endpoint, different message based on detection results ✅ Clean Implementation: - No additional endpoints needed - Uses existing /detect/hardcoded endpoint - Simple conditional message logic - Maintains original test count and structure Now the test will show the appropriate message whether memory modules are detected or not, using the same hardcoded test image.
369 lines
12 KiB
Python
369 lines
12 KiB
Python
#!/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)
|