Aherobo Ovie Victor
b54da61121
✅ Enhanced API Documentation Access: - Added /docs route to main app with instructions for Swagger UI - Created helpful documentation page with setup instructions - Added API Documentation button to web interface - Updated /api endpoint to include Swagger UI information ✅ User-Friendly Documentation: - Clear step-by-step instructions to access Swagger UI - Direct link to Swagger UI (when running) - Quick API reference on docs page - Professional styling for documentation page ✅ Improved Navigation: - Added 'API Documentation' button to main interface - Opens in new tab for easy reference - Back link to main interface - Clear visual hierarchy and instructions Now users can easily access API documentation from the main interface
414 lines
14 KiB
Python
414 lines
14 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)',
|
|
'/docs': 'GET - API documentation (Swagger UI)',
|
|
'/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),
|
|
'swagger_ui': 'http://localhost:5003/docs/ (run: python3 api_docs.py)'
|
|
})
|
|
|
|
@app.route('/docs')
|
|
def api_docs():
|
|
"""Redirect to API documentation."""
|
|
return """
|
|
<!DOCTYPE html>
|
|
<html>
|
|
<head>
|
|
<title>API Documentation</title>
|
|
<style>
|
|
body { font-family: Arial, sans-serif; text-align: center; padding: 50px; background: #f5f5f5; }
|
|
.container { max-width: 600px; margin: 0 auto; background: white; padding: 40px; border-radius: 10px; box-shadow: 0 4px 6px rgba(0,0,0,0.1); }
|
|
.btn { display: inline-block; padding: 12px 24px; background: #007bff; color: white; text-decoration: none; border-radius: 5px; margin: 10px; }
|
|
.btn:hover { background: #0056b3; }
|
|
.code { background: #f8f9fa; padding: 10px; border-radius: 5px; margin: 20px 0; font-family: monospace; }
|
|
</style>
|
|
</head>
|
|
<body>
|
|
<div class="container">
|
|
<h1>🚀 Memory Module Detection API Documentation</h1>
|
|
<p>Interactive Swagger UI documentation for all API endpoints</p>
|
|
|
|
<h3>📖 Access Swagger UI:</h3>
|
|
<div class="code">
|
|
<strong>Step 1:</strong> Start API docs server<br>
|
|
<code>python3 api_docs.py</code><br><br>
|
|
<strong>Step 2:</strong> Open Swagger UI<br>
|
|
<code>http://localhost:5003/docs/</code>
|
|
</div>
|
|
|
|
<a href="http://localhost:5003/docs/" class="btn" target="_blank">
|
|
📚 Open Swagger UI (if running)
|
|
</a>
|
|
|
|
<h3>📋 Quick API Reference:</h3>
|
|
<ul style="text-align: left;">
|
|
<li><strong>POST /detect</strong> - Upload image for detection</li>
|
|
<li><strong>GET /detect/hardcoded</strong> - Test with predefined image</li>
|
|
<li><strong>POST /detect/base64</strong> - Process base64 encoded image</li>
|
|
<li><strong>GET /health</strong> - System health check</li>
|
|
</ul>
|
|
|
|
<p><a href="/">← Back to Main Interface</a></p>
|
|
</div>
|
|
</body>
|
|
</html>
|
|
"""
|
|
|
|
@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)
|