initial commit
This commit is contained in:
+57
@@ -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
|
||||||
@@ -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
|
||||||
|
@@ -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)
|
||||||
@@ -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
|
||||||
@@ -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)}")
|
||||||
@@ -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
|
||||||
@@ -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()
|
||||||
@@ -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>
|
||||||
@@ -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)}")
|
||||||
@@ -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
|
||||||
@@ -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/<filename>`
|
||||||
|
- 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 |
|
||||||
@@ -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.
|
||||||
Reference in New Issue
Block a user