264 lines
10 KiB
Python
264 lines
10 KiB
Python
|
|
#!/usr/bin/env python3
|
|||
|
|
"""
|
|||
|
|
Swagger UI Documentation Server for Memory Module Detection API
|
|||
|
|
This server provides interactive API documentation similar to Mini SpecsComply Pro.
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import io
|
|||
|
|
import base64
|
|||
|
|
from flask import Flask, request, jsonify
|
|||
|
|
from flask_cors import CORS
|
|||
|
|
from flask_restx import Api, Resource, fields, reqparse
|
|||
|
|
from PIL import Image
|
|||
|
|
import numpy as np
|
|||
|
|
from werkzeug.utils import secure_filename
|
|||
|
|
from werkzeug.datastructures import FileStorage
|
|||
|
|
import tempfile
|
|||
|
|
import logging
|
|||
|
|
import time
|
|||
|
|
from datetime import datetime
|
|||
|
|
from inference_utils import MemoryModuleDetector
|
|||
|
|
|
|||
|
|
# Configure logging
|
|||
|
|
logging.basicConfig(level=logging.INFO)
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
# Initialize Flask app for Swagger UI only
|
|||
|
|
app = Flask(__name__)
|
|||
|
|
CORS(app)
|
|||
|
|
|
|||
|
|
# Initialize Flask-RESTX API with professional styling
|
|||
|
|
api = Api(
|
|||
|
|
app,
|
|||
|
|
version='1.0.0',
|
|||
|
|
title='Memory Module Detection API',
|
|||
|
|
description='''
|
|||
|
|
🔍 **AI-Powered Memory Module Detection System**
|
|||
|
|
|
|||
|
|
Professional computer vision API for detecting memory modules in motherboard images using YOLOv8.
|
|||
|
|
|
|||
|
|
**Features:**
|
|||
|
|
- Real-time memory module detection
|
|||
|
|
- 99.5% accuracy with YOLOv8 Nano
|
|||
|
|
- Multiple input formats (upload, base64, hardcoded test)
|
|||
|
|
- Confidence threshold control
|
|||
|
|
- Annotated image output
|
|||
|
|
|
|||
|
|
**Use Cases:**
|
|||
|
|
- Electronic waste recycling facilities
|
|||
|
|
- Hardware inventory management
|
|||
|
|
- Quality control in manufacturing
|
|||
|
|
- Educational computer vision projects
|
|||
|
|
''',
|
|||
|
|
doc='/',
|
|||
|
|
prefix='/api/v1',
|
|||
|
|
contact='Memory Module Detection Team',
|
|||
|
|
contact_email='support@memorydetection.ai'
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# Configuration
|
|||
|
|
UPLOAD_FOLDER = 'uploads'
|
|||
|
|
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'bmp'}
|
|||
|
|
MODEL_PATH = 'runs/detect/memory_module_detection/weights/best.pt'
|
|||
|
|
HARDCODED_IMAGE_PATH = 'training/memory/out1.png'
|
|||
|
|
|
|||
|
|
# Initialize detector
|
|||
|
|
detector = MemoryModuleDetector(MODEL_PATH)
|
|||
|
|
|
|||
|
|
# Create upload folder
|
|||
|
|
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
|
|||
|
|
|
|||
|
|
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."""
|
|||
|
|
buffered = io.BytesIO()
|
|||
|
|
image.save(buffered, format="PNG")
|
|||
|
|
img_str = base64.b64encode(buffered.getvalue()).decode()
|
|||
|
|
return img_str
|
|||
|
|
|
|||
|
|
# Create namespaces with descriptions
|
|||
|
|
ns_health = api.namespace('health', description='🏥 Health Check Operations')
|
|||
|
|
ns_detection = api.namespace('detection', description='🔍 Memory Module Detection Operations')
|
|||
|
|
ns_info = api.namespace('info', description='ℹ️ API Information')
|
|||
|
|
|
|||
|
|
# Define comprehensive API models
|
|||
|
|
detection_bbox = api.model('DetectionBBox', {
|
|||
|
|
'x1': fields.Float(required=True, description='Top-left X coordinate'),
|
|||
|
|
'y1': fields.Float(required=True, description='Top-left Y coordinate'),
|
|||
|
|
'x2': fields.Float(required=True, description='Bottom-right X coordinate'),
|
|||
|
|
'y2': fields.Float(required=True, description='Bottom-right Y coordinate'),
|
|||
|
|
'confidence': fields.Float(required=True, description='Detection confidence score'),
|
|||
|
|
'class': fields.String(required=True, description='Detected class name')
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
detection_result = api.model('DetectionResult', {
|
|||
|
|
'success': fields.Boolean(required=True, description='Whether detection was successful'),
|
|||
|
|
'detections': fields.List(fields.Nested(detection_bbox), description='List of detected memory modules with coordinates'),
|
|||
|
|
'num_detections': fields.Integer(description='Number of memory modules detected'),
|
|||
|
|
'annotated_image': fields.String(description='Base64 encoded annotated image with bounding boxes'),
|
|||
|
|
'confidence_threshold': fields.Float(description='Confidence threshold used for detection'),
|
|||
|
|
'test_image_path': fields.String(description='Path to the test image (for hardcoded tests)'),
|
|||
|
|
'processing_time_ms': fields.Float(description='Processing time in milliseconds')
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
error_response = api.model('ErrorResponse', {
|
|||
|
|
'error': fields.String(required=True, description='Detailed error message'),
|
|||
|
|
'success': fields.Boolean(required=True, description='Always false for errors'),
|
|||
|
|
'error_code': fields.String(description='Error classification code')
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
health_response = api.model('HealthResponse', {
|
|||
|
|
'status': fields.String(required=True, description='Overall health status', enum=['healthy', 'degraded', 'unhealthy']),
|
|||
|
|
'model_loaded': fields.Boolean(required=True, description='Whether the YOLOv8 model is loaded'),
|
|||
|
|
'model_path': fields.String(description='Path to the loaded model'),
|
|||
|
|
'timestamp': fields.String(required=True, description='Current timestamp in ISO format'),
|
|||
|
|
'uptime_seconds': fields.Float(description='API uptime in seconds'),
|
|||
|
|
'memory_usage_mb': fields.Float(description='Current memory usage in MB')
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
model_info = api.model('ModelInfo', {
|
|||
|
|
'architecture': fields.String(description='Model architecture name'),
|
|||
|
|
'version': fields.String(description='Model version'),
|
|||
|
|
'classes': fields.List(fields.String, description='Detectable object classes'),
|
|||
|
|
'input_size': fields.String(description='Expected input image size'),
|
|||
|
|
'model_loaded': fields.Boolean(description='Model loading status'),
|
|||
|
|
'accuracy_metrics': fields.Raw(description='Model performance metrics')
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
api_info_response = api.model('ApiInfoResponse', {
|
|||
|
|
'name': fields.String(required=True, description='API service name'),
|
|||
|
|
'version': fields.String(required=True, description='API version'),
|
|||
|
|
'description': fields.String(required=True, description='API description'),
|
|||
|
|
'model_info': fields.Nested(model_info, description='Information about the ML model'),
|
|||
|
|
'endpoints': fields.List(fields.String, description='Available API endpoints'),
|
|||
|
|
'supported_formats': fields.List(fields.String, description='Supported image formats'),
|
|||
|
|
'max_file_size': fields.String(description='Maximum file upload size'),
|
|||
|
|
'rate_limits': fields.Raw(description='API rate limiting information')
|
|||
|
|
})
|
|||
|
|
|
|||
|
|
# ============================================================================
|
|||
|
|
# API RESOURCES
|
|||
|
|
# ============================================================================
|
|||
|
|
|
|||
|
|
@ns_health.route('')
|
|||
|
|
class HealthCheck(Resource):
|
|||
|
|
@ns_health.doc('health_check')
|
|||
|
|
@ns_health.marshal_with(health_response)
|
|||
|
|
def get(self):
|
|||
|
|
"""
|
|||
|
|
🏥 **Check API Health Status**
|
|||
|
|
|
|||
|
|
Returns comprehensive health information including model status, uptime, and system metrics.
|
|||
|
|
Use this endpoint to monitor API availability and performance.
|
|||
|
|
"""
|
|||
|
|
import psutil
|
|||
|
|
import time
|
|||
|
|
|
|||
|
|
return {
|
|||
|
|
'status': 'healthy' if detector.model is not None else 'degraded',
|
|||
|
|
'model_loaded': detector.model is not None,
|
|||
|
|
'model_path': MODEL_PATH,
|
|||
|
|
'timestamp': datetime.now().isoformat(),
|
|||
|
|
'uptime_seconds': time.time() - start_time,
|
|||
|
|
'memory_usage_mb': psutil.Process().memory_info().rss / 1024 / 1024
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
@ns_info.route('')
|
|||
|
|
class ApiInfo(Resource):
|
|||
|
|
@ns_info.doc('api_info')
|
|||
|
|
@ns_info.marshal_with(api_info_response)
|
|||
|
|
def get(self):
|
|||
|
|
"""
|
|||
|
|
ℹ️ **Get Comprehensive API Information**
|
|||
|
|
|
|||
|
|
Returns detailed information about the API capabilities, model specifications,
|
|||
|
|
supported formats, and available endpoints.
|
|||
|
|
"""
|
|||
|
|
return {
|
|||
|
|
'name': 'Memory Module Detection API',
|
|||
|
|
'version': '1.0.0',
|
|||
|
|
'description': 'AI-powered memory module detection system for motherboard images using YOLOv8',
|
|||
|
|
'model_info': {
|
|||
|
|
'architecture': 'YOLOv8 Nano',
|
|||
|
|
'version': '8.0.196',
|
|||
|
|
'classes': ['memory_module'],
|
|||
|
|
'input_size': '640x640',
|
|||
|
|
'model_loaded': detector.model is not None,
|
|||
|
|
'accuracy_metrics': {
|
|||
|
|
'mAP50': 0.995,
|
|||
|
|
'precision': 1.0,
|
|||
|
|
'recall': 0.984,
|
|||
|
|
'inference_time_ms': 37
|
|||
|
|
}
|
|||
|
|
},
|
|||
|
|
'endpoints': [
|
|||
|
|
'/api/v1/health',
|
|||
|
|
'/api/v1/info',
|
|||
|
|
'/api/v1/detection/upload',
|
|||
|
|
'/api/v1/detection/hardcoded',
|
|||
|
|
'/api/v1/detection/base64'
|
|||
|
|
],
|
|||
|
|
'supported_formats': ['PNG', 'JPG', 'JPEG', 'GIF', 'BMP'],
|
|||
|
|
'max_file_size': '16MB',
|
|||
|
|
'rate_limits': {
|
|||
|
|
'requests_per_minute': 60,
|
|||
|
|
'concurrent_requests': 10
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
# File upload parser with detailed documentation
|
|||
|
|
upload_parser = reqparse.RequestParser()
|
|||
|
|
upload_parser.add_argument(
|
|||
|
|
'file',
|
|||
|
|
location='files',
|
|||
|
|
type=FileStorage,
|
|||
|
|
required=True,
|
|||
|
|
help='📁 Image file containing motherboard to analyze (PNG, JPG, JPEG, GIF, BMP)'
|
|||
|
|
)
|
|||
|
|
upload_parser.add_argument(
|
|||
|
|
'confidence',
|
|||
|
|
type=float,
|
|||
|
|
default=0.8,
|
|||
|
|
help='🎯 Confidence threshold for detection (0.0-1.0, default: 0.8)'
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
@ns_detection.route('/upload')
|
|||
|
|
class DetectionUpload(Resource):
|
|||
|
|
@ns_detection.doc('upload_detection')
|
|||
|
|
@ns_detection.expect(upload_parser)
|
|||
|
|
@ns_detection.marshal_with(detection_result, code=200)
|
|||
|
|
@ns_detection.marshal_with(error_response, code=400)
|
|||
|
|
@ns_detection.marshal_with(error_response, code=500)
|
|||
|
|
def post(self):
|
|||
|
|
"""
|
|||
|
|
📤 **Upload Image for Memory Module Detection**
|
|||
|
|
|
|||
|
|
Upload a motherboard image and get real-time memory module detection results.
|
|||
|
|
|
|||
|
|
**Process:**
|
|||
|
|
1. Upload image file (max 16MB)
|
|||
|
|
2. AI processes image with YOLOv8
|
|||
|
|
3. Returns detected memory modules with bounding boxes
|
|||
|
|
4. Includes annotated image with visual markers
|
|||
|
|
|
|||
|
|
**Supported Formats:** PNG, JPG, JPEG, GIF, BMP
|
|||
|
|
"""
|
|||
|
|
# This endpoint connects to the main API
|
|||
|
|
return {'error': 'This is documentation only. Use the main API at port 5002', 'success': False}, 501
|
|||
|
|
|
|||
|
|
# Global start time for uptime calculation
|
|||
|
|
start_time = time.time()
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
import time
|
|||
|
|
print("🚀 Starting Memory Module Detection API Documentation Server...")
|
|||
|
|
print("📚 Swagger UI available at: http://localhost:5003/")
|
|||
|
|
print("🔗 Main API running at: http://localhost:5002")
|
|||
|
|
print("📖 Interactive documentation with professional interface")
|
|||
|
|
|
|||
|
|
app.run(host='0.0.0.0', port=5003, debug=True)
|