154 lines
4.4 KiB
Python
154 lines
4.4 KiB
Python
import os
|
|
from flask import Flask, render_template, send_from_directory, jsonify, request
|
|
|
|
app = Flask(__name__)
|
|
|
|
IMAGES_DIR = 'images'
|
|
LABELS_DIR = 'labels'
|
|
os.makedirs(LABELS_DIR, exist_ok=True)
|
|
|
|
ALLOWED_EXTENSIONS = {'.png', '.jpg', '.jpeg', '.bmp', '.webp'}
|
|
|
|
def get_image_files():
|
|
files = []
|
|
if os.path.exists(IMAGES_DIR):
|
|
for f in os.listdir(IMAGES_DIR):
|
|
_, ext = os.path.splitext(f)
|
|
if ext.lower() in ALLOWED_EXTENSIONS:
|
|
files.append(f)
|
|
return sorted(files)
|
|
|
|
def is_labeled(filename):
|
|
name, _ = os.path.splitext(filename)
|
|
label_path = os.path.join(LABELS_DIR, name + '.txt')
|
|
return os.path.exists(label_path)
|
|
|
|
def get_yolo_boxes(filename):
|
|
name, _ = os.path.splitext(filename)
|
|
label_path = os.path.join(LABELS_DIR, name + '.txt')
|
|
boxes = []
|
|
if os.path.exists(label_path):
|
|
with open(label_path, 'r') as f:
|
|
for line in f:
|
|
parts = line.strip().split()
|
|
if len(parts) >= 5:
|
|
# format: class x_center y_center width height
|
|
try:
|
|
boxes.append({
|
|
'class': parts[0],
|
|
'x': float(parts[1]),
|
|
'y': float(parts[2]),
|
|
'w': float(parts[3]),
|
|
'h': float(parts[4])
|
|
})
|
|
except ValueError:
|
|
continue
|
|
return boxes
|
|
|
|
@app.route('/')
|
|
def index():
|
|
return render_template('index.html')
|
|
|
|
@app.route('/images/<path:filename>')
|
|
def serve_image(filename):
|
|
return send_from_directory(IMAGES_DIR, filename)
|
|
|
|
@app.route('/api/next')
|
|
def get_next_image():
|
|
images = get_image_files()
|
|
total = len(images)
|
|
labeled_count = 0
|
|
next_index = None
|
|
next_image = None
|
|
|
|
for i, img in enumerate(images):
|
|
if is_labeled(img):
|
|
labeled_count += 1
|
|
elif next_image is None:
|
|
next_image = img
|
|
next_index = i
|
|
|
|
# If all labeled, next_image remains None
|
|
|
|
return jsonify({
|
|
'total': total,
|
|
'labeled_count': labeled_count,
|
|
'index': next_index if next_index is not None else -1,
|
|
'filename': next_image,
|
|
'completed': next_image is None,
|
|
'boxes': get_yolo_boxes(next_image) if next_image else []
|
|
})
|
|
|
|
@app.route('/api/image/<int:index>')
|
|
def get_image_by_index(index):
|
|
images = get_image_files()
|
|
total = len(images)
|
|
|
|
if total == 0:
|
|
return jsonify({'error': 'No images'}), 404
|
|
|
|
# Clamp index
|
|
if index < 0: index = 0
|
|
if index >= total: index = total - 1
|
|
|
|
filename = images[index]
|
|
|
|
# Calculate stats
|
|
labeled_count = sum(1 for img in images if is_labeled(img))
|
|
|
|
return jsonify({
|
|
'total': total,
|
|
'labeled_count': labeled_count,
|
|
'index': index,
|
|
'filename': filename,
|
|
'labeled': is_labeled(filename),
|
|
'boxes': get_yolo_boxes(filename)
|
|
})
|
|
|
|
@app.route('/api/save', methods=['POST'])
|
|
def save_label():
|
|
data = request.json
|
|
filename = data.get('filename')
|
|
boxes = data.get('boxes', []) # List of {x, y, w, h} in image pixels
|
|
img_w = data.get('width')
|
|
img_h = data.get('height')
|
|
|
|
if not filename or not img_w or not img_h:
|
|
return jsonify({'error': 'Missing data'}), 400
|
|
|
|
name, _ = os.path.splitext(filename)
|
|
label_path = os.path.join(LABELS_DIR, name + '.txt')
|
|
|
|
lines = []
|
|
for box in boxes:
|
|
# YOLO format: class x_center y_center width height (normalized 0-1)
|
|
# Class is 0 for signature
|
|
|
|
# Ensure box is valid
|
|
x = float(box['x'])
|
|
y = float(box['y'])
|
|
w = float(box['w'])
|
|
h = float(box['h'])
|
|
|
|
# Center coordinates
|
|
center_x = (x + w / 2.0) / img_w
|
|
center_y = (y + h / 2.0) / img_h
|
|
norm_w = w / img_w
|
|
norm_h = h / img_h
|
|
|
|
# Clamp values just in case
|
|
center_x = max(0.0, min(1.0, center_x))
|
|
center_y = max(0.0, min(1.0, center_y))
|
|
norm_w = max(0.0, min(1.0, norm_w))
|
|
norm_h = max(0.0, min(1.0, norm_h))
|
|
|
|
lines.append(f"0 {center_x:.6f} {center_y:.6f} {norm_w:.6f} {norm_h:.6f}")
|
|
|
|
with open(label_path, 'w') as f:
|
|
f.write('\n'.join(lines))
|
|
|
|
return jsonify({'success': True})
|
|
|
|
if __name__ == '__main__':
|
|
app.run(debug=True, port=8080)
|