From db057c746757916c1361ca419fd422d9a204fdc5 Mon Sep 17 00:00:00 2001 From: Ayomide Date: Thu, 17 Jul 2025 00:03:03 +0100 Subject: [PATCH] initial commit --- .gitignore | 57 ++++++++++++++ annotations.csv | 79 ++++++++++++++++++++ backend/app.py | 77 +++++++++++++++++++ backend/config.py | 12 +++ backend/convert_csv_to_yolo.py | 38 ++++++++++ backend/detector.py | 55 ++++++++++++++ backend/main.py | 79 ++++++++++++++++++++ backend/test.html | 10 +++ backend/train.py | 132 +++++++++++++++++++++++++++++++++ detection/utils.py | 19 +++++ docs/API_Documentation.md | 75 +++++++++++++++++++ README.md => docs/README.md | 0 docs/Task.md | 26 +++++++ 13 files changed, 659 insertions(+) create mode 100644 .gitignore create mode 100644 annotations.csv create mode 100644 backend/app.py create mode 100644 backend/config.py create mode 100644 backend/convert_csv_to_yolo.py create mode 100644 backend/detector.py create mode 100644 backend/main.py create mode 100644 backend/test.html create mode 100644 backend/train.py create mode 100644 detection/utils.py create mode 100644 docs/API_Documentation.md rename README.md => docs/README.md (100%) create mode 100644 docs/Task.md diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..b8297b0 --- /dev/null +++ b/.gitignore @@ -0,0 +1,57 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# Virtual environment +venv/ +.env/ +.venv/ + +# IDE specific files +.vscode/ +.idea/ +*.swp +*.swo + +# Training artifacts +runs/ +*.pt +*.onnx +*.engine + +# Dataset files +training/ +yolo_dataset/ +*.jpg +*.png +*.txt + +# API runtime files +static/uploads/ +static/results/ +*.jpg +*.jpeg +*.png + +# Log files +*.log +logs/ + +# System files +.DS_Store +Thumbs.db + +# Python packaging +*.egg-info/ +dist/ +build/ + +# Local configuration +config.yml +local_settings.py + +# Large files +*.zip +*.tar.gz +*.pth \ No newline at end of file diff --git a/annotations.csv b/annotations.csv new file mode 100644 index 0000000..924a6b4 --- /dev/null +++ b/annotations.csv @@ -0,0 +1,79 @@ +class,x1,y1,x2,y2,filename,img_width,img_height +memory_module,541,567,661,265,out1.png,1920,1080 +memory_module,623,585,726,279,out1.png,1920,1080 +memory_module,1063,277,1055,599,out1.png,1920,1080 +memory_module,1124,287,1137,603,out1.png,1920,1080 +memory_module,533,519,833,265,out10.png,1920,1080 +memory_module,591,552,879,294,out10.png,1920,1080 +memory_module,926,717,1190,390,out10.png,1920,1080 +memory_module,997,759,1261,394,out10.png,1920,1080 +memory_module,530,517,854,262,out11.png,1920,1080 +memory_module,573,545,904,287,out11.png,1920,1080 +memory_module,544,497,890,265,out12.png,1920,1080 +memory_module,591,520,926,279,out12.png,1920,1080 +memory_module,880,724,1222,404,out12.png,1920,1080 +memory_module,955,774,1286,429,out12.png,1920,1080 +memory_module,548,470,897,251,out13.png,1920,1080 +memory_module,580,506,933,279,out13.png,1920,1080 +memory_module,851,717,1211,415,out13.png,1920,1080 +memory_module,923,767,1282,444,out13.png,1920,1080 +memory_module,562,453,922,251,out14.png,1920,1080 +memory_module,594,495,954,276,out14.png,1920,1080 +memory_module,848,702,1218,422,out14.png,1920,1080 +memory_module,905,759,1282,461,out14.png,1920,1080 +memory_module,905,767,1314,486,out15.png,1920,1080 +memory_module,858,710,1265,437,out15.png,1920,1080 +memory_module,633,478,1011,265,out15.png,1920,1080 +memory_module,591,442,961,251,out15.png,1920,1080 +memory_module,576,424,975,247,out16.png,1920,1080 +memory_module,605,456,1015,272,out16.png,1920,1080 +memory_module,819,710,1275,451,out16.png,1920,1080 +memory_module,873,763,1325,501,out16.png,1920,1080 +memory_module,555,410,968,244,out17.png,1920,1080 +memory_module,576,463,1011,279,out17.png,1920,1080 +memory_module,762,717,1261,472,out17.png,1920,1080 +memory_module,790,784,1311,515,out17.png,1920,1080 +memory_module,541,420,983,272,out18.png,1920,1080 +memory_module,576,460,1015,294,out18.png,1920,1080 +memory_module,716,752,1257,515,out18.png,1920,1080 +memory_module,755,820,1304,579,out18.png,1920,1080 +memory_module,744,834,1347,608,out19.png,1920,1080 +memory_module,726,759,1300,561,out19.png,1920,1080 +memory_module,1068,322,611,451,out19.png,1920,1080 +memory_module,598,410,1040,287,out19.png,1920,1080 +memory_module,555,581,669,269,out2.png,1920,1080 +memory_module,626,588,726,276,out2.png,1920,1080 +memory_module,1057,605,1066,285,out2.png,1920,1080 +memory_module,1135,612,1124,285,out2.png,1920,1080 +memory_module,626,395,1086,287,out20.png,1920,1080 +memory_module,641,435,1100,326,out20.png,1920,1080 +memory_module,741,752,1307,586,out20.png,1920,1080 +memory_module,748,813,1350,629,out20.png,1920,1080 +memory_module,541,595,676,279,out3.png,1920,1080 +memory_module,619,599,733,294,out3.png,1920,1080 +memory_module,1040,613,1058,292,out3.png,1920,1080 +memory_module,1122,617,1122,297,out3.png,1920,1080 +memory_module,548,602,690,297,out4.png,1920,1080 +memory_module,619,613,747,294,out4.png,1920,1080 +memory_module,1042,630,1072,305,out4.png,1920,1080 +memory_module,1122,635,1135,305,out4.png,1920,1080 +memory_module,498,620,672,312,out5.png,1920,1080 +memory_module,573,624,726,312,out5.png,1920,1080 +memory_module,1003,665,1058,319,out5.png,1920,1080 +memory_module,1086,678,1120,322,out5.png,1920,1080 +memory_module,487,624,690,326,out6.png,1920,1080 +memory_module,569,645,740,308,out6.png,1920,1080 +memory_module,975,696,1076,345,out6.png,1920,1080 +memory_module,1061,728,1138,353,out6.png,1920,1080 +memory_module,1021,751,1153,370,out7.png,1920,1080 +memory_module,938,724,1089,357,out7.png,1920,1080 +memory_module,533,632,761,311,out7.png,1920,1080 +memory_module,468,612,708,314,out7.png,1920,1080 +memory_module,490,594,750,312,out8.png,1920,1080 +memory_module,553,615,798,297,out8.png,1920,1080 +memory_module,933,739,1126,379,out8.png,1920,1080 +memory_module,1011,764,1195,381,out8.png,1920,1080 +memory_module,523,574,797,294,out9.png,1920,1080 +memory_module,583,595,847,287,out9.png,1920,1080 +memory_module,955,742,1168,383,out9.png,1920,1080 +memory_module,1030,767,1232,397,out9.png,1920,1080 \ No newline at end of file diff --git a/backend/app.py b/backend/app.py new file mode 100644 index 0000000..efe9134 --- /dev/null +++ b/backend/app.py @@ -0,0 +1,77 @@ +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 + +app = Flask(__name__) +logging.basicConfig(level=logging.INFO) + +# Initialize detector +MODEL_PATH = str(Path(__file__).parent.parent / "runs" / "detect" / "train" / "weights" / "best.pt") +model = YOLO(MODEL_PATH) + + +@app.route('/') +def index(): + return send_file('test.html') + + +@app.route('/detect', methods=['POST']) +def detect(): + if 'image' not in request.files: + return jsonify({'error': 'No image provided'}), 400 + + file = request.files['image'] + if file.filename == '': + return jsonify({'error': 'No selected file'}), 400 + + 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) + + # Generate annotated image + annotated = results[0].plot(line_width=2, font_size=0.5) + + # Save to results folder + output_dir = Path("static/results") + output_dir.mkdir(exist_ok=True) + 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]) + }) + + return jsonify({ + 'detections': detections, + 'result_image': f"/results/{filename}" + }) + + except Exception as e: + logging.error(f"Detection error: {str(e)}") + return jsonify({'error': str(e)}), 500 + +@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 diff --git a/backend/config.py b/backend/config.py new file mode 100644 index 0000000..4fca4c1 --- /dev/null +++ b/backend/config.py @@ -0,0 +1,12 @@ +import os + + +class Config: + UPLOAD_FOLDER = 'static/uploads' + RESULT_FOLDER = 'static/results' + HARDCODED_IMAGES = 'training/memory' + ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'} + + # Model configuration + MODEL_PATH = 'models/memory_detector.pt' + CONFIDENCE_THRESHOLD = 0.5 \ No newline at end of file diff --git a/backend/convert_csv_to_yolo.py b/backend/convert_csv_to_yolo.py new file mode 100644 index 0000000..73ede7f --- /dev/null +++ b/backend/convert_csv_to_yolo.py @@ -0,0 +1,38 @@ +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 new file mode 100644 index 0000000..1a1d54f --- /dev/null +++ b/backend/detector.py @@ -0,0 +1,55 @@ +from ultralytics import YOLO +import cv2 +from pathlib import Path +import logging + +logger = logging.getLogger(__name__) + + +class MemoryDetector: + def __init__(self, model_path): + try: + self.model = YOLO(model_path) + logger.info(f"Loaded model from {model_path}") + except Exception as e: + logger.error(f"Model loading failed: {str(e)}") + raise + + def detect(self, image_path): + try: + # Run inference + results = self.model.predict(image_path, imgsz=416, conf=0.5) + + # Extract results + boxes = results[0].boxes.xyxy.cpu().numpy() + confidences = results[0].boxes.conf.cpu().numpy() + + # 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) + }) + + # Annotate image + annotated_img = self._draw_boxes(image_path, detections) + + return { + 'detections': detections, + 'annotated_image': annotated_img + } + + 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 diff --git a/backend/main.py b/backend/main.py new file mode 100644 index 0000000..3443580 --- /dev/null +++ b/backend/main.py @@ -0,0 +1,79 @@ +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/test.html b/backend/test.html new file mode 100644 index 0000000..964562e --- /dev/null +++ b/backend/test.html @@ -0,0 +1,10 @@ + + + +

Memory Module Detection

+
+ + +
+ + \ No newline at end of file diff --git a/backend/train.py b/backend/train.py new file mode 100644 index 0000000..a422e44 --- /dev/null +++ b/backend/train.py @@ -0,0 +1,132 @@ +import os +import yaml +from pathlib import Path +from ultralytics import YOLO +import logging +from sklearn.model_selection import train_test_split +import torch + +# Setup logging +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + + +class DatasetPreparer: + def __init__(self): + # Get the project root + self.project_root = Path(__file__).parent.parent + self.training_dir = self.project_root / "training" + self.output_dir = self.project_root / "yolo_dataset" + + logger.info(f"Looking for training data in: {self.training_dir}") + logger.info(f"Output will be saved to: {self.output_dir}") + + def verify_dataset(self): + """Check if images exist in the correct structure""" + memory_images = list((self.training_dir / "memory").glob("*.[jJ][pP][gG]")) + \ + list((self.training_dir / "memory").glob("*.[pP][nN][gG]")) + no_memory_images = list((self.training_dir / "no_memory").glob("*.[jJ][pP][gG]")) + \ + list((self.training_dir / "no_memory").glob("*.[pP][nN][gG]")) + + if not memory_images: + raise FileNotFoundError(f"No images found in {self.training_dir/'memory/'}") + if not no_memory_images: + logger.warning(f"No images found in {self.training_dir/'no_memory/'}") + + logger.info(f"Found {len(memory_images)} memory images and {len(no_memory_images)} no_memory images") + return memory_images + no_memory_images + + def organize_yolo_dataset(self, test_size=0.2): + """Organize into YOLO directory structure""" + try: + all_images = self.verify_dataset() + + # Create directories + (self.output_dir / "images/train").mkdir(parents=True, exist_ok=True) + (self.output_dir / "images/val").mkdir(parents=True, exist_ok=True) + (self.output_dir / "labels/train").mkdir(parents=True, exist_ok=True) + (self.output_dir / "labels/val").mkdir(parents=True, exist_ok=True) + + # Split into train/val + train_files, val_files = train_test_split(all_images, test_size=test_size, random_state=42) + + # Create symlinks (or copy files) + for file in train_files: + dest = self.output_dir / "images/train" / file.name + if not dest.exists(): + os.link(str(file), str(dest)) + + # Handle annotations if they exist + label_file = file.with_suffix('.txt') + if label_file.exists(): + label_dest = self.output_dir / "labels/train" / label_file.name + if not label_dest.exists(): + os.link(str(label_file), str(label_dest)) + + for file in val_files: + dest = self.output_dir / "images/val" / file.name + if not dest.exists(): + os.link(str(file), str(dest)) + + label_file = file.with_suffix('.txt') + if label_file.exists(): + label_dest = self.output_dir / "labels/val" / label_file.name + if not label_dest.exists(): + os.link(str(label_file), str(label_dest)) + + # Create dataset YAML + data = { + 'train': str(self.output_dir / "images/train"), + 'val': str(self.output_dir / "images/val"), + 'nc': 1, + 'names': ['memory_module'] + } + + with open(self.output_dir / "dataset.yaml", 'w') as f: + yaml.dump(data, f) + + logger.info("YOLO dataset prepared successfully") + return True + + except Exception as e: + logger.error(f"Error organizing dataset: {str(e)}") + return False + + +def train_model(): + """Train YOLO model using ultralytics""" + try: + model = YOLO('yolov8n.pt') + + results = model.train( + data=str(Path(__file__).parent.parent / "yolo_dataset/dataset.yaml"), + epochs=100, # Reduced from 300 for local testing + batch=2, # Small batch size for limited VRAM + imgsz=416, # Reduced from 640 to save memory + device='0' if torch.cuda.is_available() else 'cpu', + augment=True, # for small datasets + 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 + cache=False, # Disable cache if low on disk space + single_cls=True, + optimizer='AdamW', # For small datasets + seed=42, + pretrained=True # Using pretrained weights + ) + + logger.info("Training completed successfully") + return True + except Exception as e: + logger.error(f"Training failed: {str(e)}") + return False + + +if __name__ == "__main__": + try: + preparer = DatasetPreparer() + if preparer.organize_yolo_dataset(): + train_model() + except Exception as e: + logger.error(f"Fatal error: {str(e)}") diff --git a/detection/utils.py b/detection/utils.py new file mode 100644 index 0000000..f46abf7 --- /dev/null +++ b/detection/utils.py @@ -0,0 +1,19 @@ +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 new file mode 100644 index 0000000..aedc976 --- /dev/null +++ b/docs/API_Documentation.md @@ -0,0 +1,75 @@ +# Memory Module Detection API Documentation + +## Overview +Flask API for detecting memory modules on motherboard images using YOLOv8. Processes uploaded images and returns bounding box coordinates with confidence scores. + +## Base URL +`http://localhost:5000` + +## Endpoints + +### 1. Root Endpoint +**GET** `/` +- Returns the test interface HTML page +- Response: `test.html` + +### 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 + +### 3. Result Retrieval +**GET** `/results/` +- Returns annotated image with bounding boxes +- Example: `http://localhost:5000/results/out1.jpg` + +## Request/Response Examples +**Sample Request:** +```python +import requests +response = requests.post( + 'http://localhost:5000/detect', + files={'image': open('motherboard.jpg', 'rb')} +) +``` + +**Sample Response:** +```json +{ + "detections": [ + { + "box": [541,567,661,265], + "confidence": 0.98, + "class": 0 + } + ], + "result_image": "/results/out1.jpg" +} +``` + +## Technical Specifications +| Parameter | Value | +|--------------------|---------------------------| +| Model | YOLOv8n (custom-trained) | +| Input Formats | JPG/PNG | +| Recommended Resolution | 416px | +| Processing Time (CPU) | 200-500ms per image | diff --git a/README.md b/docs/README.md similarity index 100% rename from README.md rename to docs/README.md diff --git a/docs/Task.md b/docs/Task.md new file mode 100644 index 0000000..c855caa --- /dev/null +++ b/docs/Task.md @@ -0,0 +1,26 @@ +**1. Algorithm Choice** +- **Selected:** YOLOv8n (lightweight version) +- **Why:** + - Fast detection (0.5s/image on CPU) + - Works well with small datasets (40 images) + - Accurate for motherboard components + +**2. Hardware Impact** +- **Training:** + - GPU recommended (4x faster training) + - CPU works but slower +- **Deployment:** + - CPU sufficient for basic use + - GPU better for high volume + +**3. Video Handling** +- **Approach:** Process each frame individually +- **Changes Needed:** + - Add frame-by-frame processing + - Include tracking to follow memory modules + - Optimize for speed (lower resolution helps) + +**Key Facts:** +- Same model works for images/video +- CPU processing is practical +- No architecture changes needed between image/video modes. \ No newline at end of file