Ayomide
77 lines
2.2 KiB
Python
77 lines
2.2 KiB
Python
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) |