update
This commit is contained in:
Aherobo Ovie Victor
@@ -0,0 +1,14 @@
|
||||
from flask import Flask
|
||||
|
||||
def create_app():
|
||||
app = Flask(__name__)
|
||||
|
||||
# Register blueprints
|
||||
from app.routes import main_bp
|
||||
app.register_blueprint(main_bp)
|
||||
|
||||
# Ensure the static folder is properly set
|
||||
app.static_folder = 'static'
|
||||
app.template_folder = 'templates'
|
||||
|
||||
return app
|
||||
@@ -0,0 +1,57 @@
|
||||
from flask import Blueprint, request, jsonify, send_file, render_template
|
||||
from app.utils.detector import MemoryDetector
|
||||
import os
|
||||
from PIL import Image
|
||||
import io
|
||||
|
||||
main_bp = Blueprint('main', __name__)
|
||||
detector = MemoryDetector()
|
||||
|
||||
@main_bp.route('/')
|
||||
def index():
|
||||
return render_template('index.html')
|
||||
|
||||
@main_bp.route('/detect', methods=['POST'])
|
||||
def detect_memory():
|
||||
if 'image' not in request.files:
|
||||
return jsonify({'error': 'No image provided'}), 400
|
||||
|
||||
file = request.files['image']
|
||||
|
||||
# Read the image
|
||||
img = Image.open(file.stream)
|
||||
|
||||
# Process the image and get annotated image and detections
|
||||
annotated_img, detections = detector.detect(img)
|
||||
|
||||
# Convert PIL image to bytes
|
||||
img_byte_arr = io.BytesIO()
|
||||
annotated_img.save(img_byte_arr, format='PNG')
|
||||
img_byte_arr.seek(0)
|
||||
|
||||
return send_file(
|
||||
img_byte_arr,
|
||||
mimetype='image/png'
|
||||
)
|
||||
|
||||
@main_bp.route('/detect/test', methods=['GET'])
|
||||
def detect_test():
|
||||
"""Endpoint for testing with a hardcoded image"""
|
||||
# Using an existing image from the validation set
|
||||
test_image_path = os.path.join('training', 'val', 'images', 'memory_out19.png')
|
||||
|
||||
if not os.path.exists(test_image_path):
|
||||
return jsonify({'error': f'Test image not found at {test_image_path}'}), 404
|
||||
|
||||
img = Image.open(test_image_path)
|
||||
# Get both the annotated image and detections
|
||||
annotated_img, detections = detector.detect(img)
|
||||
|
||||
img_byte_arr = io.BytesIO()
|
||||
annotated_img.save(img_byte_arr, format='PNG')
|
||||
img_byte_arr.seek(0)
|
||||
|
||||
return send_file(
|
||||
img_byte_arr,
|
||||
mimetype='image/png'
|
||||
)
|
||||
@@ -0,0 +1,159 @@
|
||||
* {
|
||||
margin: 0;
|
||||
padding: 0;
|
||||
box-sizing: border-box;
|
||||
}
|
||||
|
||||
body {
|
||||
font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell, sans-serif;
|
||||
line-height: 1.6;
|
||||
background-color: #f5f5f5;
|
||||
color: #333;
|
||||
}
|
||||
|
||||
.container {
|
||||
max-width: 1200px;
|
||||
margin: 0 auto;
|
||||
padding: 2rem;
|
||||
}
|
||||
|
||||
h1 {
|
||||
text-align: center;
|
||||
margin-bottom: 2rem;
|
||||
color: #2c3e50;
|
||||
}
|
||||
|
||||
.upload-section {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
align-items: center;
|
||||
gap: 1rem;
|
||||
margin-bottom: 2rem;
|
||||
}
|
||||
|
||||
.upload-box {
|
||||
width: 100%;
|
||||
max-width: 500px;
|
||||
height: 200px;
|
||||
border: 2px dashed #3498db;
|
||||
border-radius: 8px;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
cursor: pointer;
|
||||
transition: border-color 0.3s ease;
|
||||
background-color: #fff;
|
||||
}
|
||||
|
||||
.upload-box:hover {
|
||||
border-color: #2980b9;
|
||||
}
|
||||
|
||||
.upload-content {
|
||||
text-align: center;
|
||||
}
|
||||
|
||||
.upload-icon {
|
||||
width: 64px;
|
||||
height: 64px;
|
||||
margin-bottom: 1rem;
|
||||
}
|
||||
|
||||
.browse-text {
|
||||
color: #3498db;
|
||||
text-decoration: underline;
|
||||
cursor: pointer;
|
||||
}
|
||||
|
||||
button {
|
||||
padding: 0.8rem 1.5rem;
|
||||
font-size: 1rem;
|
||||
border: none;
|
||||
border-radius: 4px;
|
||||
cursor: pointer;
|
||||
transition: background-color 0.3s ease;
|
||||
}
|
||||
|
||||
#detectButton {
|
||||
background-color: #3498db;
|
||||
color: white;
|
||||
}
|
||||
|
||||
#detectButton:disabled {
|
||||
background-color: #bdc3c7;
|
||||
cursor: not-allowed;
|
||||
}
|
||||
|
||||
#testButton {
|
||||
background-color: #2ecc71;
|
||||
color: white;
|
||||
}
|
||||
|
||||
#detectButton:hover:not(:disabled) {
|
||||
background-color: #2980b9;
|
||||
}
|
||||
|
||||
#testButton:hover {
|
||||
background-color: #27ae60;
|
||||
}
|
||||
|
||||
.results-section {
|
||||
margin-top: 2rem;
|
||||
}
|
||||
|
||||
.image-container {
|
||||
display: flex;
|
||||
gap: 2rem;
|
||||
justify-content: center;
|
||||
flex-wrap: wrap;
|
||||
}
|
||||
|
||||
.image-box {
|
||||
flex: 1;
|
||||
min-width: 300px;
|
||||
max-width: 500px;
|
||||
background-color: white;
|
||||
padding: 1rem;
|
||||
border-radius: 8px;
|
||||
box-shadow: 0 2px 4px rgba(0,0,0,0.1);
|
||||
}
|
||||
|
||||
.image-box h3 {
|
||||
margin-bottom: 1rem;
|
||||
text-align: center;
|
||||
}
|
||||
|
||||
.image-box img {
|
||||
width: 100%;
|
||||
height: auto;
|
||||
border-radius: 4px;
|
||||
display: none;
|
||||
}
|
||||
|
||||
.loading-spinner {
|
||||
position: fixed;
|
||||
top: 0;
|
||||
left: 0;
|
||||
width: 100%;
|
||||
height: 100%;
|
||||
background-color: rgba(255,255,255,0.8);
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
z-index: 1000;
|
||||
}
|
||||
|
||||
.spinner {
|
||||
width: 50px;
|
||||
height: 50px;
|
||||
border: 5px solid #f3f3f3;
|
||||
border-top: 5px solid #3498db;
|
||||
border-radius: 50%;
|
||||
animation: spin 1s linear infinite;
|
||||
}
|
||||
|
||||
@keyframes spin {
|
||||
0% { transform: rotate(0deg); }
|
||||
100% { transform: rotate(360deg); }
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 2.9 MiB |
@@ -0,0 +1,5 @@
|
||||
<svg xmlns="http://www.w3.org/2000/svg" width="64" height="64" viewBox="0 0 24 24" fill="none" stroke="#3498db" stroke-width="2" stroke-linecap="round" stroke-linejoin="round">
|
||||
<path d="M21 15v4a2 2 0 0 1-2 2H5a2 2 0 0 1-2-2v-4"/>
|
||||
<polyline points="17 8 12 3 7 8"/>
|
||||
<line x1="12" y1="3" x2="12" y2="15"/>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 323 B |
@@ -0,0 +1,100 @@
|
||||
document.addEventListener('DOMContentLoaded', function() {
|
||||
const dropZone = document.getElementById('dropZone');
|
||||
const fileInput = document.getElementById('fileInput');
|
||||
const detectButton = document.getElementById('detectButton');
|
||||
const testButton = document.getElementById('testButton');
|
||||
const originalImage = document.getElementById('originalImage');
|
||||
const resultImage = document.getElementById('resultImage');
|
||||
const loading = document.getElementById('loading');
|
||||
|
||||
// Handle drag and drop
|
||||
dropZone.addEventListener('dragover', (e) => {
|
||||
e.preventDefault();
|
||||
dropZone.style.borderColor = '#2980b9';
|
||||
});
|
||||
|
||||
dropZone.addEventListener('dragleave', (e) => {
|
||||
e.preventDefault();
|
||||
dropZone.style.borderColor = '#3498db';
|
||||
});
|
||||
|
||||
dropZone.addEventListener('drop', (e) => {
|
||||
e.preventDefault();
|
||||
dropZone.style.borderColor = '#3498db';
|
||||
|
||||
const file = e.dataTransfer.files[0];
|
||||
if (file && file.type.startsWith('image/')) {
|
||||
handleImageSelection(file);
|
||||
}
|
||||
});
|
||||
|
||||
// Handle click to upload
|
||||
dropZone.addEventListener('click', () => {
|
||||
fileInput.click();
|
||||
});
|
||||
|
||||
fileInput.addEventListener('change', (e) => {
|
||||
const file = e.target.files[0];
|
||||
if (file) {
|
||||
handleImageSelection(file);
|
||||
}
|
||||
});
|
||||
|
||||
function handleImageSelection(file) {
|
||||
const reader = new FileReader();
|
||||
reader.onload = function(e) {
|
||||
originalImage.src = e.target.result;
|
||||
originalImage.style.display = 'block';
|
||||
detectButton.disabled = false;
|
||||
};
|
||||
reader.readAsDataURL(file);
|
||||
}
|
||||
|
||||
// Handle detect button click
|
||||
detectButton.addEventListener('click', async () => {
|
||||
const formData = new FormData();
|
||||
formData.append('image', fileInput.files[0]);
|
||||
|
||||
try {
|
||||
loading.style.display = 'flex';
|
||||
const response = await fetch('/detect', {
|
||||
method: 'POST',
|
||||
body: formData
|
||||
});
|
||||
|
||||
if (response.ok) {
|
||||
const blob = await response.blob();
|
||||
resultImage.src = URL.createObjectURL(blob);
|
||||
resultImage.style.display = 'block';
|
||||
} else {
|
||||
alert('Error processing image');
|
||||
}
|
||||
} catch (error) {
|
||||
console.error('Error:', error);
|
||||
alert('Error processing image');
|
||||
} finally {
|
||||
loading.style.display = 'none';
|
||||
}
|
||||
});
|
||||
|
||||
// Handle test button click
|
||||
testButton.addEventListener('click', async () => {
|
||||
try {
|
||||
loading.style.display = 'flex';
|
||||
const response = await fetch('/detect/test');
|
||||
|
||||
if (response.ok) {
|
||||
const blob = await response.blob();
|
||||
resultImage.src = URL.createObjectURL(blob);
|
||||
resultImage.style.display = 'block';
|
||||
} else {
|
||||
alert('Error running test detection');
|
||||
}
|
||||
} catch (error) {
|
||||
console.error('Error:', error);
|
||||
alert('Error running test detection');
|
||||
} finally {
|
||||
loading.style.display = 'none';
|
||||
}
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,47 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>Memory Module Detector</title>
|
||||
<link rel="stylesheet" href="{{ url_for('static', filename='css/style.css') }}">
|
||||
</head>
|
||||
<body>
|
||||
<div class="container">
|
||||
<h1>Memory Module Detector</h1>
|
||||
|
||||
<div class="upload-section">
|
||||
<div class="upload-box" id="dropZone">
|
||||
<input type="file" id="fileInput" accept="image/*" hidden>
|
||||
<div class="upload-content">
|
||||
<img src="{{ url_for('static', filename='images/upload-icon.svg') }}" alt="Upload" class="upload-icon">
|
||||
<p>Drag and drop an image or <span class="browse-text">browse</span></p>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<button id="detectButton" disabled>Detect Memory Modules</button>
|
||||
<button id="testButton">Run Test Detection</button>
|
||||
</div>
|
||||
|
||||
<div class="results-section">
|
||||
<div class="image-container">
|
||||
<div class="image-box">
|
||||
<h3>Original Image</h3>
|
||||
<img id="originalImage" src="" alt="Original image will appear here">
|
||||
</div>
|
||||
<div class="image-box">
|
||||
<h3>Detected Results</h3>
|
||||
<img id="resultImage" src="" alt="Detection results will appear here">
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div id="loading" class="loading-spinner" style="display: none;">
|
||||
<div class="spinner"></div>
|
||||
<p>Processing...</p>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<script src="{{ url_for('static', filename='js/main.js') }}"></script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -0,0 +1,99 @@
|
||||
from ultralytics import YOLO
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
from typing import Tuple, List, Dict
|
||||
|
||||
class MemoryDetector:
|
||||
def __init__(self,
|
||||
model_path='model/weights/best.pt',
|
||||
conf_threshold=0.25,
|
||||
iou_threshold=0.45):
|
||||
"""
|
||||
Initialize the detector with the trained model.
|
||||
|
||||
Args:
|
||||
model_path (str): Path to the trained model weights
|
||||
conf_threshold (float): Confidence threshold for detections
|
||||
iou_threshold (float): IoU threshold for NMS
|
||||
"""
|
||||
self.model = YOLO(model_path)
|
||||
self.conf_threshold = conf_threshold
|
||||
self.iou_threshold = iou_threshold
|
||||
|
||||
def detect(self,
|
||||
image: Image.Image,
|
||||
conf_threshold: float = None,
|
||||
iou_threshold: float = None) -> Tuple[Image.Image, List[Dict]]:
|
||||
"""
|
||||
Detect memory modules in the given image.
|
||||
|
||||
Args:
|
||||
image (PIL.Image): Input image to process
|
||||
conf_threshold (float, optional): Override default confidence threshold
|
||||
iou_threshold (float, optional): Override default IoU threshold
|
||||
|
||||
Returns:
|
||||
Tuple[PIL.Image, List[Dict]]: Annotated image and list of detections
|
||||
"""
|
||||
# Use provided thresholds or defaults
|
||||
conf = conf_threshold if conf_threshold is not None else self.conf_threshold
|
||||
iou = iou_threshold if iou_threshold is not None else self.iou_threshold
|
||||
|
||||
# Run inference
|
||||
results = self.model.predict(
|
||||
source=image,
|
||||
conf=conf,
|
||||
iou=iou,
|
||||
max_det=10,
|
||||
verbose=False
|
||||
)
|
||||
|
||||
# Get the annotated image
|
||||
annotated_img = results[0].plot()
|
||||
|
||||
# Extract detection information
|
||||
detections = []
|
||||
for box in results[0].boxes:
|
||||
detection = {
|
||||
'xyxy': box.xyxy[0].tolist(), # Bounding box coordinates
|
||||
'confidence': float(box.conf[0]), # Detection confidence
|
||||
'class': int(box.cls[0]) # Class ID
|
||||
}
|
||||
detections.append(detection)
|
||||
|
||||
return Image.fromarray(annotated_img), detections
|
||||
|
||||
def optimize_thresholds(self, validation_images: List[Image.Image]) -> Tuple[float, float]:
|
||||
"""
|
||||
Find optimal confidence and IoU thresholds using validation images.
|
||||
|
||||
Args:
|
||||
validation_images (List[Image.Image]): List of validation images
|
||||
|
||||
Returns:
|
||||
Tuple[float, float]: Optimal confidence and IoU thresholds
|
||||
"""
|
||||
best_conf = 0.25
|
||||
best_iou = 0.45
|
||||
|
||||
# Grid search for best parameters
|
||||
conf_range = [0.15, 0.2, 0.25, 0.3, 0.35]
|
||||
iou_range = [0.35, 0.4, 0.45, 0.5, 0.55]
|
||||
|
||||
best_score = 0
|
||||
|
||||
for conf in conf_range:
|
||||
for iou in iou_range:
|
||||
total_score = 0
|
||||
for img in validation_images:
|
||||
_, detections = self.detect(img, conf, iou)
|
||||
# Score based on number of detections and confidence
|
||||
score = sum([d['confidence'] for d in detections])
|
||||
total_score += score
|
||||
|
||||
if total_score > best_score:
|
||||
best_score = total_score
|
||||
best_conf = conf
|
||||
best_iou = iou
|
||||
|
||||
return best_conf, best_iou
|
||||
Reference in New Issue
Block a user