update project structure and improve scripts
This commit is contained in:
Ayomide
+211
-51
@@ -1,77 +1,237 @@
|
||||
from flask import Flask, request, jsonify, send_file
|
||||
from pathlib import Path
|
||||
from ultralytics import YOLO
|
||||
import cv2
|
||||
import numpy as np
|
||||
import os
|
||||
import uuid
|
||||
import logging
|
||||
from io import BytesIO
|
||||
import random
|
||||
from pathlib import Path
|
||||
from flask import Flask, request, jsonify, send_file, render_template
|
||||
import cv2
|
||||
|
||||
from config import Config
|
||||
from detector import MemoryDetector
|
||||
from exceptions import DetectionError, FileUploadError, ValidationError, VideoProcessingError
|
||||
from utils import allowed_file, setup_logging
|
||||
from video_processor import VideoProcessor
|
||||
|
||||
# Initialize Flask app
|
||||
app = Flask(__name__)
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
app.config['MAX_CONTENT_LENGTH'] = Config.MAX_FILE_SIZE
|
||||
|
||||
# Initialize detector
|
||||
MODEL_PATH = str(Path(__file__).parent.parent / "runs" / "detect" / "train" / "weights" / "best.pt")
|
||||
model = YOLO(MODEL_PATH)
|
||||
# Setup logging
|
||||
setup_logging()
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Initialize detector and video processor
|
||||
detector = None
|
||||
video_processor = None
|
||||
|
||||
|
||||
def init_app():
|
||||
"""Initialize application."""
|
||||
global detector, video_processor
|
||||
|
||||
try:
|
||||
# Create directories
|
||||
Config.create_directories()
|
||||
|
||||
# Validate model exists
|
||||
Config.validate_model()
|
||||
|
||||
# Initialize detector
|
||||
detector = MemoryDetector(
|
||||
model_path=Config.MODEL_PATH,
|
||||
confidence_threshold=Config.CONFIDENCE_THRESHOLD,
|
||||
image_size=Config.IMAGE_SIZE
|
||||
)
|
||||
|
||||
# Initialize video processor
|
||||
video_processor = VideoProcessor(detector)
|
||||
|
||||
logger.info("Application initialized successfully")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Application initialization failed: {e}")
|
||||
raise
|
||||
|
||||
|
||||
@app.route('/')
|
||||
def index():
|
||||
return send_file('test.html')
|
||||
"""Serve the frontend interface."""
|
||||
return render_template('index.html')
|
||||
|
||||
|
||||
@app.route('/detect', methods=['POST'])
|
||||
def detect():
|
||||
if 'image' not in request.files:
|
||||
return jsonify({'error': 'No image provided'}), 400
|
||||
@app.route('/api/health')
|
||||
def health_check():
|
||||
"""Health check endpoint."""
|
||||
return jsonify({
|
||||
'status': 'healthy',
|
||||
'service': 'Memory Detection API',
|
||||
'model_loaded': detector is not None
|
||||
})
|
||||
|
||||
file = request.files['image']
|
||||
if file.filename == '':
|
||||
return jsonify({'error': 'No selected file'}), 400
|
||||
|
||||
@app.route('/api/v1/detect', methods=['POST'])
|
||||
def detect_memory():
|
||||
"""Detect memory modules in uploaded image."""
|
||||
try:
|
||||
# Read image directly from memory
|
||||
img_bytes = file.read()
|
||||
img = cv2.imdecode(np.frombuffer(img_bytes, np.uint8), cv2.IMREAD_COLOR)
|
||||
|
||||
# Run prediction
|
||||
results = model.predict(img, imgsz=416, conf=0.3)
|
||||
# Validate file upload
|
||||
if 'image' not in request.files:
|
||||
raise FileUploadError("No image file provided")
|
||||
|
||||
# Generate annotated image
|
||||
annotated = results[0].plot(line_width=2, font_size=0.5)
|
||||
file = request.files['image']
|
||||
if file.filename == '':
|
||||
raise FileUploadError("No file selected")
|
||||
|
||||
# Save to results folder
|
||||
output_dir = Path("static/results")
|
||||
output_dir.mkdir(exist_ok=True)
|
||||
if not allowed_file(file.filename, Config.ALLOWED_EXTENSIONS):
|
||||
raise ValidationError(f"Invalid file type. Allowed: {Config.ALLOWED_EXTENSIONS}")
|
||||
|
||||
# Process image
|
||||
logger.info(f"Processing uploaded file: {file.filename}")
|
||||
|
||||
# Read image data
|
||||
image_bytes = file.read()
|
||||
|
||||
# Perform detection
|
||||
result = detector.detect_from_bytes(image_bytes)
|
||||
|
||||
# Save result image
|
||||
filename = f"{uuid.uuid4()}.jpg"
|
||||
output_path = output_dir / filename
|
||||
cv2.imwrite(str(output_path), annotated)
|
||||
|
||||
# Extract detections
|
||||
detections = []
|
||||
for box in results[0].boxes:
|
||||
detections.append({
|
||||
'box': box.xyxy[0].tolist(),
|
||||
'confidence': float(box.conf[0]),
|
||||
'class': int(box.cls[0])
|
||||
})
|
||||
output_path = Path(Config.RESULT_FOLDER) / filename
|
||||
cv2.imwrite(str(output_path), result['annotated_image'])
|
||||
|
||||
return jsonify({
|
||||
'detections': detections,
|
||||
'result_image': f"/results/{filename}"
|
||||
'success': True,
|
||||
'data': {
|
||||
'detections': result['detections'],
|
||||
'detection_count': result['detection_count'],
|
||||
'result_image_url': f"/api/v1/results/{filename}"
|
||||
}
|
||||
})
|
||||
|
||||
except (DetectionError, FileUploadError, ValidationError) as e:
|
||||
logger.error(f"Detection request failed: {e}")
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': str(e)
|
||||
}), 400
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error in detection: {e}")
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': 'Internal server error'
|
||||
}), 500
|
||||
|
||||
|
||||
@app.route('/api/v1/results/<filename>')
|
||||
def get_result_image(filename):
|
||||
"""Get result image."""
|
||||
try:
|
||||
file_path = Path(Config.RESULT_FOLDER) / filename
|
||||
if not file_path.exists():
|
||||
return jsonify({'error': 'Image not found'}), 404
|
||||
|
||||
return send_file(file_path)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error serving result image: {e}")
|
||||
return jsonify({'error': 'Internal server error'}), 500
|
||||
|
||||
|
||||
@app.route('/api/v1/detect/video', methods=['POST'])
|
||||
def detect_video():
|
||||
"""Process video for memory module detection."""
|
||||
try:
|
||||
# Validate video upload
|
||||
if 'video' not in request.files:
|
||||
raise FileUploadError("No video file provided")
|
||||
|
||||
file = request.files['video']
|
||||
if file.filename == '':
|
||||
raise FileUploadError("No file selected")
|
||||
|
||||
# Check video format
|
||||
video_extensions = {'mp4', 'avi', 'mov', 'mkv'}
|
||||
if not allowed_file(file.filename, video_extensions):
|
||||
raise ValidationError(f"Invalid video format. Allowed: {video_extensions}")
|
||||
|
||||
# Save uploaded video temporarily
|
||||
video_filename = f"temp_{uuid.uuid4()}.mp4"
|
||||
video_path = Path(Config.UPLOAD_FOLDER) / video_filename
|
||||
file.save(str(video_path))
|
||||
|
||||
# Get processing parameters
|
||||
fps = request.form.get('fps', 1, type=int)
|
||||
max_frames = request.form.get('max_frames', 50, type=int)
|
||||
|
||||
logger.info(f"Processing video: {file.filename}, fps={fps}, max_frames={max_frames}")
|
||||
|
||||
# Process video
|
||||
result = video_processor.process_video(
|
||||
str(video_path),
|
||||
fps=fps,
|
||||
max_frames=max_frames
|
||||
)
|
||||
|
||||
# Clean up temporary file
|
||||
video_path.unlink(missing_ok=True)
|
||||
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'data': {
|
||||
'video_filename': file.filename,
|
||||
'processing_info': result['processing_info'],
|
||||
'detections': result['detections'],
|
||||
'summary': result['summary']
|
||||
}
|
||||
})
|
||||
|
||||
except (VideoProcessingError, FileUploadError, ValidationError) as e:
|
||||
logger.error(f"Video processing failed: {e}")
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': str(e)
|
||||
}), 400
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error in video processing: {e}")
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': 'Internal server error'
|
||||
}), 500
|
||||
|
||||
|
||||
@app.route('/api/v1/test-images')
|
||||
def list_test_images():
|
||||
"""List available test images."""
|
||||
try:
|
||||
test_images = Config.get_test_images()
|
||||
return jsonify({
|
||||
'success': True,
|
||||
'test_images': [img.name for img in test_images]
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Detection error: {str(e)}")
|
||||
return jsonify({'error': str(e)}), 500
|
||||
logger.error(f"Error listing test images: {e}")
|
||||
return jsonify({'error': 'Internal server error'}), 500
|
||||
|
||||
|
||||
@app.errorhandler(413)
|
||||
def file_too_large(e):
|
||||
"""Handle file too large error."""
|
||||
return jsonify({
|
||||
'success': False,
|
||||
'error': f'File too large. Maximum size: {Config.MAX_FILE_SIZE} bytes'
|
||||
}), 413
|
||||
|
||||
@app.route('/results/<filename>')
|
||||
def get_result(filename):
|
||||
return send_file(Path("static/results") / filename)
|
||||
|
||||
if __name__ == '__main__':
|
||||
# Create directories
|
||||
Path("static/uploads").mkdir(parents=True, exist_ok=True)
|
||||
Path("static/results").mkdir(parents=True, exist_ok=True)
|
||||
app.run(host='0.0.0.0', port=5000)
|
||||
# Initialize application
|
||||
init_app()
|
||||
|
||||
# Start server
|
||||
logger.info(f"Starting server on {Config.HOST}:{Config.PORT}")
|
||||
app.run(
|
||||
host=Config.HOST,
|
||||
port=Config.PORT,
|
||||
debug=Config.DEBUG
|
||||
)
|
||||
Reference in New Issue
Block a user