| import os |
| import traceback |
| import uuid |
|
|
| import torch |
| from dotenv import load_dotenv |
| from flask import Flask, jsonify, render_template, request |
| from flask_cors import CORS |
| from ultralytics import YOLO |
| from werkzeug.utils import secure_filename |
|
|
| from processing import process_images |
|
|
|
|
| BASE_DIR = os.path.dirname(os.path.abspath(__file__)) |
| load_dotenv(os.path.join(BASE_DIR, '.env')) |
|
|
| app = Flask(__name__) |
| CORS(app) |
|
|
| ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'} |
|
|
|
|
| def resolve_app_path(configured_path, default_relative_path): |
| """Resolve configured relative paths from the application directory.""" |
| path = configured_path or default_relative_path |
| if os.path.isabs(path): |
| return path |
| return os.path.join(BASE_DIR, path) |
|
|
|
|
| UPLOAD_FOLDER = resolve_app_path(os.getenv('UPLOAD_FOLDER'), os.path.join('static', 'uploads')) |
| MODELS_FOLDER = resolve_app_path(os.getenv('MODELS_FOLDER'), 'models') |
| PARTS_MODEL_NAME = os.getenv('PARTS_MODEL_NAME', 'best_parts_EP336.pt') |
| DAMAGE_MODEL_NAME = os.getenv('DAMAGE_MODEL_NAME', 'best_new_EP382.pt') |
|
|
|
|
| def resolve_model_path(model_name): |
| """Resolve an absolute model path or a filename relative to MODELS_FOLDER.""" |
| if os.path.isabs(model_name): |
| return model_name |
| return os.path.join(MODELS_FOLDER, model_name) |
|
|
|
|
| PARTS_MODEL_PATH = resolve_model_path(PARTS_MODEL_NAME) |
| DAMAGE_MODEL_PATH = resolve_model_path(DAMAGE_MODEL_NAME) |
|
|
| app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER |
| os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True) |
| os.makedirs(MODELS_FOLDER, exist_ok=True) |
|
|
| device = 'cuda' if torch.cuda.is_available() else 'cpu' |
| print(f'Using device: {device}') |
|
|
|
|
| def load_model(model_path, label): |
| """Load a YOLO model if its configured file exists.""" |
| if not os.path.isfile(model_path): |
| print(f'Warning: {label} model file not found at {model_path}') |
| return None |
|
|
| try: |
| model = YOLO(model_path) |
| model.to(device) |
| print(f"Successfully loaded {label.lower()} model '{os.path.basename(model_path)}' on {device}.") |
| return model |
| except Exception as exc: |
| print(f'Error loading {label} model ({model_path}): {exc}') |
| traceback.print_exc() |
| return None |
|
|
|
|
| parts_model = load_model(PARTS_MODEL_PATH, 'Parts') |
| damage_model = load_model(DAMAGE_MODEL_PATH, 'Damage') |
|
|
|
|
| def allowed_file(filename): |
| """Return True when filename has an allowed image extension.""" |
| return bool(filename) and '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS |
|
|
|
|
| def model_readiness_error(): |
| """Return a useful model readiness payload, or None when both models are loaded.""" |
| missing = [] |
| if parts_model is None: |
| missing.append({'model': 'parts', 'path': PARTS_MODEL_PATH}) |
| if damage_model is None: |
| missing.append({'model': 'damage', 'path': DAMAGE_MODEL_PATH}) |
|
|
| if not missing: |
| return None |
|
|
| return { |
| 'error': 'Prediction models are not ready.', |
| 'missing_or_unloaded_models': missing, |
| 'hint': 'Place the configured .pt files at the listed paths or update MODELS_FOLDER/model names in .env.' |
| } |
|
|
|
|
| @app.route('/') |
| def home(): |
| """Serve the main HTML page.""" |
| return render_template('index.html') |
|
|
|
|
| @app.route('/health', methods=['GET']) |
| def health(): |
| """Expose basic service/model readiness without running inference.""" |
| readiness = model_readiness_error() |
| if readiness: |
| return jsonify({'status': 'degraded', **readiness}), 503 |
| return jsonify({'status': 'ok', 'device': device}), 200 |
|
|
|
|
| @app.route('/predict', methods=['POST']) |
| def predict(): |
| """Receive one or more images and return part/damage predictions.""" |
| readiness = model_readiness_error() |
| if readiness: |
| return jsonify(readiness), 503 |
|
|
| raw_session_key = request.form.get('session_key', '').strip() |
| session_key = secure_filename(raw_session_key)[:100] or uuid.uuid4().hex |
|
|
| if 'file' not in request.files: |
| return jsonify({'error': 'No file part in the request'}), 400 |
|
|
| files = request.files.getlist('file') |
| if not files or all(not file.filename for file in files): |
| return jsonify({'error': 'No selected files'}), 400 |
|
|
| saved_images = [] |
| skipped_files = [] |
|
|
| try: |
| for index, file in enumerate(files): |
| original_filename = file.filename or '' |
| if not (file and allowed_file(original_filename)): |
| skipped_files.append(original_filename or f'file_{index}') |
| continue |
|
|
| safe_original = secure_filename(original_filename) or f'image_{index}.jpg' |
| unique_filename = f'{uuid.uuid4().hex}_{safe_original}' |
| filepath = os.path.join(app.config['UPLOAD_FOLDER'], unique_filename) |
| file.save(filepath) |
| saved_images.append({ |
| 'path': filepath, |
| 'filename': original_filename, |
| 'index': index, |
| }) |
|
|
| if not saved_images: |
| return jsonify({ |
| 'error': 'No valid files were uploaded. Allowed types: png, jpg, jpeg', |
| 'skipped_files': skipped_files, |
| }), 400 |
|
|
| print(f"Processing {len(saved_images)} file(s) for session '{session_key}'...") |
| results = process_images(parts_model, damage_model, saved_images) |
| print(f"Processing complete for session '{session_key}'.") |
| return jsonify(results) |
|
|
| except Exception as exc: |
| print(f'An error occurred during processing for session {session_key}: {exc}') |
| traceback.print_exc() |
| return jsonify({'error': f'An error occurred during processing: {exc}'}), 500 |
|
|
| finally: |
| for image in saved_images: |
| filepath = image['path'] |
| try: |
| if os.path.exists(filepath): |
| os.remove(filepath) |
| except Exception as exc: |
| print(f'Error cleaning up file {filepath}: {exc}') |
|
|
|
|
| if __name__ == '__main__': |
| app.run(host='0.0.0.0', port=7860, debug=False) |
|
|
