From b23314375c7a70ce817d15097139780d258f7519 Mon Sep 17 00:00:00 2001 From: Ayomide Date: Thu, 24 Jul 2025 23:31:47 +0100 Subject: [PATCH] update project structure and improve scripts --- .gitignore | 8 +- backend/app.py | 262 ++++++++++++++---- backend/config.py | 59 +++- backend/convert_csv_to_yolo.py | 38 --- backend/detector.py | 151 ++++++++--- backend/exceptions.py | 31 +++ backend/main.py | 79 ------ backend/templates/index.html | 301 +++++++++++++++++++++ backend/test.html | 10 - backend/utils.py | 40 +++ backend/video_processor.py | 225 +++++++++++++++ detection/utils.py | 19 -- docs/API_Documentation.md | 195 ++++++++----- backend/train.py => scripts/train_model.py | 2 +- 14 files changed, 1117 insertions(+), 303 deletions(-) delete mode 100644 backend/convert_csv_to_yolo.py create mode 100644 backend/exceptions.py delete mode 100644 backend/main.py create mode 100644 backend/templates/index.html delete mode 100644 backend/test.html create mode 100644 backend/utils.py create mode 100644 backend/video_processor.py delete mode 100644 detection/utils.py rename backend/train.py => scripts/train_model.py (98%) diff --git a/.gitignore b/.gitignore index b8297b0..7ec0921 100644 --- a/.gitignore +++ b/.gitignore @@ -33,6 +33,7 @@ static/results/ *.jpg *.jpeg *.png +images/ # Log files *.log @@ -54,4 +55,9 @@ local_settings.py # Large files *.zip *.tar.gz -*.pth \ No newline at end of file +*.pth +*.csv +*.pt + +# Test files +tests/ \ No newline at end of file diff --git a/backend/app.py b/backend/app.py index efe9134..93db2ab 100644 --- a/backend/app.py +++ b/backend/app.py @@ -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/') +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/') -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) \ No newline at end of file + # 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 + ) \ No newline at end of file diff --git a/backend/config.py b/backend/config.py index 4fca4c1..4081b38 100644 --- a/backend/config.py +++ b/backend/config.py @@ -1,12 +1,61 @@ import os +from pathlib import Path class Config: + # Server settings + HOST = os.environ.get('HOST', '0.0.0.0') + PORT = int(os.environ.get('PORT', 5000)) + DEBUG = os.environ.get('FLASK_ENV') == 'development' + + # File settings UPLOAD_FOLDER = 'static/uploads' RESULT_FOLDER = 'static/results' - HARDCODED_IMAGES = 'training/memory' - ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'} + MAX_FILE_SIZE = 16 * 1024 * 1024 # 16MB + ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'bmp'} - # Model configuration - MODEL_PATH = 'models/memory_detector.pt' - CONFIDENCE_THRESHOLD = 0.5 \ No newline at end of file + # Model settings - pre-trained model + MODEL_PATH = os.environ.get('MODEL_PATH', 'yolov8n.pt') # Auto-downloads if not exists + CONFIDENCE_THRESHOLD = float(os.environ.get('CONFIDENCE_THRESHOLD', 0.3)) + IMAGE_SIZE = 416 + + # Hardcoded test images + TEST_IMAGES_PATH = 'training/memory' + + # Logging + LOG_LEVEL = os.environ.get('LOG_LEVEL', 'INFO') + LOG_FILE = 'logs/app.log' + + @classmethod + def create_directories(cls): + """Create required directories.""" + directories = [ + cls.UPLOAD_FOLDER, + cls.RESULT_FOLDER, + Path(cls.LOG_FILE).parent, + ] + + for directory in directories: + Path(directory).mkdir(parents=True, exist_ok=True) + + @classmethod + def validate_model(cls): + """Check if model file exists or can be downloaded.""" + # For pre-trained models + if cls.MODEL_PATH in ['yolov8n.pt', 'yolov8s.pt', 'yolov8m.pt', 'yolov8l.pt', 'yolov8x.pt']: + print(f"Using pre-trained model: {cls.MODEL_PATH} (will auto-download if needed)") + return + + # For custom models, check if file exists + if not Path(cls.MODEL_PATH).exists(): + raise FileNotFoundError(f"Model file not found: {cls.MODEL_PATH}") + + @classmethod + def get_test_images(cls): + """Get list of test images.""" + test_path = Path(cls.TEST_IMAGES_PATH) + if not test_path.exists(): + return [] + + return [f for f in test_path.iterdir() + if f.suffix.lower()[1:] in cls.ALLOWED_EXTENSIONS] diff --git a/backend/convert_csv_to_yolo.py b/backend/convert_csv_to_yolo.py deleted file mode 100644 index 73ede7f..0000000 --- a/backend/convert_csv_to_yolo.py +++ /dev/null @@ -1,38 +0,0 @@ -import pandas as pd -from pathlib import Path -import os - - -def csv_to_yolo(csv_path, output_dir): - # Create output directory if it doesn't exist - Path(output_dir).mkdir(parents=True, exist_ok=True) - - df = pd.read_csv(csv_path) - - for filename in df['filename'].unique(): - img_data = df[df['filename'] == filename].iloc[0] - img_w, img_h = img_data['img_width'], img_data['img_height'] - - yolo_lines = [] - for _, row in df[df['filename'] == filename].iterrows(): - # Convert absolute to normalized coordinates - x_center = ((row['x1'] + row['x2']) / 2) / img_w - y_center = ((row['y1'] + row['y2']) / 2) / img_h - width = abs(row['x2'] - row['x1']) / img_w - height = abs(row['y2'] - row['y1']) / img_h - - yolo_lines.append(f"0 {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}") - - # Save as YOLO .txt file - txt_path = Path(output_dir) / f"{Path(filename).stem}.txt" - with open(txt_path, 'w') as f: - f.write("\n".join(yolo_lines)) - print(f"Successfully converted CSV to YOLO format in {output_dir}") - -# Error handling -try: - csv_to_yolo("annotations.csv", "yolo_labels") -except FileNotFoundError: - print("Error: annotations.csv not found. Please check the file path.") -except Exception as e: - print(f"An error occurred: {str(e)}") \ No newline at end of file diff --git a/backend/detector.py b/backend/detector.py index 1a1d54f..e118914 100644 --- a/backend/detector.py +++ b/backend/detector.py @@ -1,55 +1,136 @@ -from ultralytics import YOLO import cv2 -from pathlib import Path +import numpy as np import logging +from pathlib import Path +from ultralytics import YOLO +from exceptions import ModelLoadError, DetectionError, ImageProcessingError logger = logging.getLogger(__name__) class MemoryDetector: - def __init__(self, model_path): + def __init__(self, model_path, confidence_threshold=0.3, image_size=416): + self.model_path = model_path + self.confidence_threshold = confidence_threshold + self.image_size = image_size + self.model = None + self._load_model() + + def _load_model(self): + """Load YOLO model with error handling.""" try: - self.model = YOLO(model_path) - logger.info(f"Loaded model from {model_path}") + if not Path(self.model_path).exists(): + raise FileNotFoundError(f"Model file not found: {self.model_path}") + + logger.info(f"Loading model from {self.model_path}") + self.model = YOLO(self.model_path) + logger.info("Model loaded successfully") + except Exception as e: - logger.error(f"Model loading failed: {str(e)}") - raise - - def detect(self, image_path): + logger.error(f"Failed to load model: {e}") + raise ModelLoadError(f"Model loading failed: {str(e)}") + + def detect_from_bytes(self, image_bytes): + """Detect memory modules from image bytes.""" try: + # Decode image + nparr = np.frombuffer(image_bytes, np.uint8) + image = cv2.imdecode(nparr, cv2.IMREAD_COLOR) + + if image is None: + raise ImageProcessingError("Could not decode image") + + return self._perform_detection(image) + + except Exception as e: + logger.error(f"Detection from bytes failed: {e}") + raise DetectionError(f"Detection failed: {str(e)}") + + def detect_from_file(self, file_path): + """Detect memory modules from image file.""" + try: + if not Path(file_path).exists(): + raise FileNotFoundError(f"Image file not found: {file_path}") + + image = cv2.imread(str(file_path)) + if image is None: + raise ImageProcessingError("Could not load image file") + + return self._perform_detection(image) + + except Exception as e: + logger.error(f"Detection from file failed: {e}") + raise DetectionError(f"Detection failed: {str(e)}") + + def _perform_detection(self, image): + """Perform detection on image.""" + try: + logger.info("Running detection") + # Run inference - results = self.model.predict(image_path, imgsz=416, conf=0.5) + results = self.model.predict( + image, + imgsz=self.image_size, + conf=self.confidence_threshold, + verbose=False + ) - # Extract results - boxes = results[0].boxes.xyxy.cpu().numpy() - confidences = results[0].boxes.conf.cpu().numpy() + # Extract detections + detections = self._extract_detections(results[0]) - # Convert to list of [x1, y1, x2, y2, confidence] - detections = [] - for box, conf in zip(boxes, confidences): - detections.append({ - 'box': [int(x) for x in box], - 'confidence': float(conf) - }) + # Create annotated image + annotated_image = self._draw_boxes(image, detections) - # Annotate image - annotated_img = self._draw_boxes(image_path, detections) + logger.info(f"Detection completed: {len(detections)} objects found") return { 'detections': detections, - 'annotated_image': annotated_img + 'annotated_image': annotated_image, + 'detection_count': len(detections) } except Exception as e: - logger.error(f"Detection failed: {str(e)}") - raise - - def _draw_boxes(self, image_path, detections): - img = cv2.imread(str(image_path)) - for det in detections: - x1, y1, x2, y2 = det['box'] - cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2) - cv2.putText(img, f"{det['confidence']:.2f}", - (x1, y1-10), - cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0), 1) - return img + logger.error(f"Detection processing failed: {e}") + raise DetectionError(f"Detection processing failed: {str(e)}") + + def _extract_detections(self, result): + """Extract detection results.""" + detections = [] + + if result.boxes is not None: + boxes = result.boxes.xyxy.cpu().numpy() + confidences = result.boxes.conf.cpu().numpy() + classes = result.boxes.cls.cpu().numpy() if result.boxes.cls is not None else None + + for i, (box, conf) in enumerate(zip(boxes, confidences)): + detection = { + 'box': [float(coord) for coord in box], # [x1, y1, x2, y2] + 'confidence': float(conf), + 'class': int(classes[i]) if classes is not None else 0 + } + detections.append(detection) + + return detections + + def _draw_boxes(self, image, detections): + """Draw bounding boxes on image.""" + annotated = image.copy() + + for detection in detections: + box = detection['box'] + confidence = detection['confidence'] + + # Extract coordinates + x1, y1, x2, y2 = map(int, box) + + # Draw bounding box + cv2.rectangle(annotated, (x1, y1), (x2, y2), (0, 255, 0), 2) + + # Draw confidence score + label = f"Memory: {confidence:.2f}" + cv2.putText( + annotated, label, (x1, y1 - 10), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1 + ) + + return annotated diff --git a/backend/exceptions.py b/backend/exceptions.py new file mode 100644 index 0000000..6b0f146 --- /dev/null +++ b/backend/exceptions.py @@ -0,0 +1,31 @@ +"""Custom exceptions for memory detection.""" + + +class DetectionError(Exception): + """Raised when detection fails.""" + pass + + +class ModelLoadError(Exception): + """Raised when model loading fails.""" + pass + + +class ImageProcessingError(Exception): + """Raised when image processing fails.""" + pass + + +class FileUploadError(Exception): + """Raised when file upload fails.""" + pass + + +class ValidationError(Exception): + """Raised when validation fails.""" + pass + + +class VideoProcessingError(Exception): + """Raised when video processing fails.""" + pass diff --git a/backend/main.py b/backend/main.py deleted file mode 100644 index 3443580..0000000 --- a/backend/main.py +++ /dev/null @@ -1,79 +0,0 @@ -import os -import sys -import logging -from pathlib import Path -from app import app -from config import Config -from training import TrainingDataManager - -# Ensuring the backend directory is in the system path -backend_dir = Path(__file__).parent -sys.path.insert(0, str(backend_dir)) - -# Configure logging -logging.basicConfig( - level=getattr(logging, Config.LOG_LEVEL), - format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' -) - -logger = logging.getLogger(__name__) - - -def setup_environment(): - """Set up the application environment.""" - try: - # Validate and create required directories - Config.validate_paths() - - # Initialize training data manager - training_manager = TrainingDataManager() - - # Log training data statistics - stats = training_manager.get_training_statistics() - logger.info(f"Training data loaded: {stats}") - - # Validate training data - validation = training_manager.validate_training_data() - logger.info(f"Training data validation: {validation}") - - return True - - except Exception as e: - logger.error(f"Error setting up environment: {e}") - return False - - -def main(): - """Main application entry point.""" - logger.info("Starting Memory Module Detection API") - - # Set up environment - if not setup_environment(): - logger.error("Failed to set up environment") - sys.exit(1) - - # Display configuration - logger.info(f"Server configuration:") - logger.info(f" - Host: {Config.HOST}") - logger.info(f" - Port: {Config.PORT}") - logger.info(f" - Debug: {Config.DEBUG}") - logger.info(f" - Algorithm: {Config.ALGORITHM}") - logger.info(f" - Max file size: {Config.MAX_CONTENT_LENGTH} bytes") - - # Start the Flask application - try: - app.run( - host=Config.HOST, - port=Config.PORT, - debug=Config.DEBUG, - threaded=True - ) - except KeyboardInterrupt: - logger.info("Application stopped by user") - except Exception as e: - logger.error(f"Application error: {e}") - sys.exit(1) - - -if __name__ == '__main__': - main() diff --git a/backend/templates/index.html b/backend/templates/index.html new file mode 100644 index 0000000..bf5d6ba --- /dev/null +++ b/backend/templates/index.html @@ -0,0 +1,301 @@ + + + + + + Memory Module Detection + + + +

Memory Module Detection

+ +
+

Health Check

+ +
+
+ +
+

Image Detection

+ + +
+
+ +
+

Video Detection

+ +
+ + + + + +
+ +
+
+ + + + \ No newline at end of file diff --git a/backend/test.html b/backend/test.html deleted file mode 100644 index 964562e..0000000 --- a/backend/test.html +++ /dev/null @@ -1,10 +0,0 @@ - - - -

Memory Module Detection

-
- - -
- - \ No newline at end of file diff --git a/backend/utils.py b/backend/utils.py new file mode 100644 index 0000000..8538b53 --- /dev/null +++ b/backend/utils.py @@ -0,0 +1,40 @@ +import os +import logging +from pathlib import Path + + +def allowed_file(filename, allowed_extensions): + """Check if file extension is allowed.""" + return '.' in filename and \ + filename.rsplit('.', 1)[1].lower() in allowed_extensions + + +def setup_logging(): + """Setup simple logging configuration.""" + # Create logs directory + log_dir = Path('logs') + log_dir.mkdir(exist_ok=True) + + # Configure logging + logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', + handlers=[ + logging.FileHandler('logs/app.log'), + logging.StreamHandler() + ] + ) + + # Set levels for noisy libraries + logging.getLogger('werkzeug').setLevel(logging.WARNING) + logging.getLogger('ultralytics').setLevel(logging.WARNING) + + +def validate_image_size(file_size, max_size): + """Validate image file size.""" + return file_size <= max_size + + +def get_file_extension(filename): + """Get file extension from filename.""" + return filename.rsplit('.', 1)[1].lower() if '.' in filename else '' \ No newline at end of file diff --git a/backend/video_processor.py b/backend/video_processor.py new file mode 100644 index 0000000..6e0c782 --- /dev/null +++ b/backend/video_processor.py @@ -0,0 +1,225 @@ +"""Simple video processing for memory module detection.""" + +import cv2 +import logging +from pathlib import Path +from typing import List, Dict, Any +from exceptions import VideoProcessingError + +logger = logging.getLogger(__name__) + + +class VideoProcessor: + """Simple video processor for memory module detection.""" + + def __init__(self, detector): + self.detector = detector + + def process_video(self, video_path: str, fps: int = 1, max_frames: int = 100) -> Dict[str, Any]: + """ + Process video file and detect memory modules in frames. + + Args: + video_path: Path to video file + fps: Frames per second to process (1 = every second) + max_frames: Maximum frames to process + + Returns: + Dictionary with detection results + """ + try: + if not Path(video_path).exists(): + raise FileNotFoundError(f"Video file not found: {video_path}") + + logger.info(f"Processing video: {video_path}") + + # Open video + cap = cv2.VideoCapture(video_path) + if not cap.isOpened(): + raise VideoProcessingError("Could not open video file") + + # Get video properties + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + video_fps = cap.get(cv2.CAP_PROP_FPS) + duration = total_frames / video_fps if video_fps > 0 else 0 + + # Calculate frame interval + frame_interval = max(1, int(video_fps / fps)) if video_fps > 0 else 1 + + logger.info(f"Video properties: {total_frames} frames, {video_fps:.2f}fps, {duration:.2f}s") + + # Process frames + detections = [] + frame_count = 0 + processed_frames = 0 + + while cap.isOpened() and processed_frames < max_frames: + ret, frame = cap.read() + if not ret: + break + + if frame_count % frame_interval == 0: + try: + # Get timestamp + timestamp = frame_count / video_fps if video_fps > 0 else frame_count + + # Detect memory modules in frame + result = self.detector._perform_detection(frame) + + # Store frame detection + frame_detection = { + 'frame_number': frame_count, + 'timestamp': round(timestamp, 2), + 'detections': result['detections'], + 'detection_count': result['detection_count'] + } + detections.append(frame_detection) + processed_frames += 1 + + logger.info(f"Frame {frame_count}: {result['detection_count']} detections") + + except Exception as e: + logger.warning(f"Failed to process frame {frame_count}: {e}") + + frame_count += 1 + + cap.release() + + # Generate summary + total_detections = sum(d['detection_count'] for d in detections) + avg_detections = total_detections / len(detections) if detections else 0 + + result = { + 'video_info': { + 'path': video_path, + 'total_frames': total_frames, + 'fps': video_fps, + 'duration': duration + }, + 'processing_info': { + 'frames_processed': processed_frames, + 'frame_interval': frame_interval, + 'target_fps': fps + }, + 'detections': detections, + 'summary': { + 'total_detections': total_detections, + 'avg_detections_per_frame': round(avg_detections, 2), + 'frames_with_detections': len([d for d in detections if d['detection_count'] > 0]) + } + } + + logger.info(f"Video processing completed: {processed_frames} frames, {total_detections} total detections") + return result + + except Exception as e: + logger.error(f"Video processing failed: {e}") + raise VideoProcessingError(f"Video processing failed: {str(e)}") + + def create_annotated_video(self, video_path: str, output_path: str, fps: int = 1) -> str: + """ + Create annotated video with bounding boxes. + + Args: + video_path: Input video path + output_path: Output video path + fps: Processing fps + + Returns: + Path to annotated video + """ + try: + logger.info(f"Creating annotated video: {video_path} -> {output_path}") + + # Open input video + cap = cv2.VideoCapture(video_path) + if not cap.isOpened(): + raise VideoProcessingError("Could not open input video") + + # Get video properties + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + video_fps = cap.get(cv2.CAP_PROP_FPS) + + # Create output video writer + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + out = cv2.VideoWriter(output_path, fourcc, video_fps, (width, height)) + + frame_interval = max(1, int(video_fps / fps)) if video_fps > 0 else 1 + frame_count = 0 + + while cap.isOpened(): + ret, frame = cap.read() + if not ret: + break + + # Process every nth frame for detection + if frame_count % frame_interval == 0: + try: + result = self.detector._perform_detection(frame) + annotated_frame = result['annotated_image'] + out.write(annotated_frame) + except Exception as e: + logger.warning(f"Failed to annotate frame {frame_count}: {e}") + out.write(frame) # Write original frame if annotation fails + else: + out.write(frame) # Write original frame + + frame_count += 1 + + cap.release() + out.release() + + logger.info(f"Annotated video created: {output_path}") + return output_path + + except Exception as e: + logger.error(f"Annotated video creation failed: {e}") + raise VideoProcessingError(f"Annotated video creation failed: {str(e)}") + + +def extract_frames(video_path: str, output_dir: str, fps: int = 1) -> List[str]: + """ + Extract frames from video for individual processing. + + Args: + video_path: Path to video file + output_dir: Directory to save frames + fps: Frames per second to extract + + Returns: + List of extracted frame paths + """ + try: + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + + cap = cv2.VideoCapture(video_path) + video_fps = cap.get(cv2.CAP_PROP_FPS) + frame_interval = max(1, int(video_fps / fps)) + + frame_paths = [] + frame_count = 0 + saved_count = 0 + + while cap.isOpened(): + ret, frame = cap.read() + if not ret: + break + + if frame_count % frame_interval == 0: + frame_filename = f"frame_{saved_count:04d}.jpg" + frame_path = output_path / frame_filename + cv2.imwrite(str(frame_path), frame) + frame_paths.append(str(frame_path)) + saved_count += 1 + + frame_count += 1 + + cap.release() + logger.info(f"Extracted {len(frame_paths)} frames to {output_dir}") + return frame_paths + + except Exception as e: + logger.error(f"Frame extraction failed: {e}") + raise VideoProcessingError(f"Frame extraction failed: {str(e)}") \ No newline at end of file diff --git a/detection/utils.py b/detection/utils.py deleted file mode 100644 index f46abf7..0000000 --- a/detection/utils.py +++ /dev/null @@ -1,19 +0,0 @@ -import cv2 - - -def draw_boxes(image, boxes): - """ - Draw bounding boxes on image - :param image: Input image - :param boxes: List of bounding boxes in format [(x1, y1, x2, y2), ...] - :return: Image with boxes drawn - """ - for box in boxes: - x1, y1, x2, y2 = box - cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2) - return image - - -def allowed_file(filename, allowed_extensions): - return '.' in filename and \ - filename.rsplit('.', 1)[1].lower() in allowed_extensions diff --git a/docs/API_Documentation.md b/docs/API_Documentation.md index aedc976..1aedfc2 100644 --- a/docs/API_Documentation.md +++ b/docs/API_Documentation.md @@ -1,75 +1,142 @@ -# Memory Module Detection API Documentation +# Memory Module Detection System -## Overview -Flask API for detecting memory modules on motherboard images using YOLOv8. Processes uploaded images and returns bounding box coordinates with confidence scores. +A computer vision system for detecting memory modules in images and videos using YOLO object detection. -## Base URL -`http://localhost:5000` +## Features -## Endpoints +- **Image Detection**: Upload images to detect memory modules +- **Video Processing**: Process video files frame-by-frame +- **Bounding Box Visualization**: See detected objects with confidence scores +- **Detailed Analytics**: Get detection statistics and metrics +- **REST API**: Fully documented API endpoints +- **Web Interface**: User-friendly browser interface -### 1. Root Endpoint -**GET** `/` -- Returns the test interface HTML page -- Response: `test.html` +## Requirements -### 2. Image Detection -**POST** `/detect` -- Accepts image uploads for processing -- **Request:** - ```bash - curl -X POST -F "image=@motherboard.jpg" http://localhost:5000/detect - ``` -- **Successful Response (200):** - ```json - { - "detections": [ - { - "box": [x1,y1,x2,y2], - "confidence": 0.95, - "class": 0 - } - ], - "result_image": "/results/filename.jpg" - } - ``` -- **Error Responses:** - - `400 Bad Request`: Missing/invalid image file - - `500 Server Error`: Processing failure +- Python 3.8+ +- CUDA-enabled GPU (recommended for best performance) +- Docker (optional) -### 3. Result Retrieval -**GET** `/results/` -- Returns annotated image with bounding boxes -- Example: `http://localhost:5000/results/out1.jpg` +## Installation -## Request/Response Examples -**Sample Request:** -```python -import requests -response = requests.post( - 'http://localhost:5000/detect', - files={'image': open('motherboard.jpg', 'rb')} -) +### 1. Clone the repository +```bash +git clone https://github.com/yourusername/memory-detection.git +cd memory-detection ``` -**Sample Response:** -```json -{ - "detections": [ - { - "box": [541,567,661,265], - "confidence": 0.98, - "class": 0 - } - ], - "result_image": "/results/out1.jpg" -} +### 2. Set up Python environment +```bash +python -m venv venv +source venv/bin/activate # Linux/MacOS +venv\Scripts\activate # Windows ``` -## Technical Specifications -| Parameter | Value | -|--------------------|---------------------------| -| Model | YOLOv8n (custom-trained) | -| Input Formats | JPG/PNG | -| Recommended Resolution | 416px | -| Processing Time (CPU) | 200-500ms per image | +### 3. Install dependencies +```bash +pip install -r requirements.txt +``` + +### 4. Download YOLO model +Place your `.pt` model file in the `models/` directory and update `config.py` with the correct path. + +## Configuration + +Edit `config.py` to customize: + +- Model paths +- Confidence thresholds +- File size limits +- Server host/port settings + +## Running the Application + +### Development Mode +```bash +python backend/app.py +``` + +### Production Mode (with Gunicorn) +```bash +gunicorn --bind 0.0.0.0:5000 backend.app:app +``` + +### Docker +```bash +docker build -t memory-detection . +docker run -p 5000:5000 memory-detection +``` + +## API Documentation + +### Endpoints + +| Endpoint | Method | Description | +|----------|--------|-------------| +| `/` | GET | Web interface | +| `/api/health` | GET | Service health check | +| `/api/v1/detect` | POST | Detect memory modules in image | +| `/api/v1/detect/video` | POST | Process video for detections | +| `/api/v1/results/` | GET | Retrieve result images | + +### Example Requests + +**Image Detection**: +```bash +curl -X POST -F "image=@test.jpg" http://localhost:5000/api/v1/detect +``` + +**Video Processing**: +```bash +curl -X POST -F "video=@test.mp4" -F "fps=2" -F "max_frames=100" http://localhost:5000/api/v1/detect/video +``` + +## Web Interface + +Access the web interface at `http://localhost:5000`: + +1. **Upload Images**: + - Supported formats: JPG, PNG + - Max file size: 10MB (configurable) + +2. **Process Videos**: + - Supported formats: MP4, AVI, MOV + - Adjustable FPS and frame limits + - Detailed frame-by-frame results + +## Project Structure + +``` +memory-detection/ +├── backend/ # Core application +│ ├── app.py # Flask application +│ ├── config.py # Configuration +│ ├── detector.py # YOLO detection logic +│ ├── exceptions.py # Custom exceptions +│ ├── utils.py # Helper functions +│ ├── video_processor.py # Video handling +│ └── templates/ # Frontend templates +│ └── index.html # Web interface +├── models/ # YOLO model files +├── uploads/ # Temporary upload storage +├── results/ # Detection result images +├── tests/ # Unit tests +├── requirements.txt # Python dependencies +└── README.md +``` + +## Troubleshooting + +**Common Issues**: + +1. **Model not loading**: + - Verify model path in `config.py` + - Check file permissions + +2. **CUDA errors**: + - Ensure compatible CUDA version is installed + - Try CPU-only mode by modifying detector.py + +3. **File upload issues**: + - Check `MAX_FILE_SIZE` in config + - Verify file extensions are allowed diff --git a/backend/train.py b/scripts/train_model.py similarity index 98% rename from backend/train.py rename to scripts/train_model.py index a422e44..6f8dc03 100644 --- a/backend/train.py +++ b/scripts/train_model.py @@ -108,7 +108,7 @@ def train_model(): patience=20, # Early stopping if no improvement lr0=0.001, # Learning rate cos_lr=True, # Cosine learning rate scheduler - workers=1, # Reduce if memory errors + workers=1, # Reduce memory errors cache=False, # Disable cache if low on disk space single_cls=True, optimizer='AdamW', # For small datasets