initial commit

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/<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)

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

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)}")

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

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()

backend/test.html

@@ -0,0 +1,10 @@
+<!DOCTYPE html>
+<html>
+<body>
+    <h2>Memory Module Detection</h2>
+    <form action="/detect" method="post" enctype="multipart/form-data">
+        <input type="file" name="image" accept="image/*">
+        <input type="submit" value="Upload">
+    </form>
+</body>
+</html>

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)}")