Нейросеть для опознавания лиц
Что нужно добавить в код, чтобы он определял лицо и пол человека.
обучение:
from ultralytics import YOLO
import subprocess
import os
import glob
import sys
bl = "run/detect/train/weight/ best.pt "
subprocess.run (f" yolo task = detect mode = train data={"dataset"}/data.yaml model = ' yolo11n.pt ' epochs = 50 imgsz = 640", shell=True)
subprocess.run (f" yolo task = detect mode = val model = {bl} data = {"dataset"}/data.yaml", shell=True)
subprocess.run (f" yolo task = detect mode = predict model = {bl} conf = 0.25 source = {"dataset"}/train/images save = True", shell=True)
Сервер:
from flask import Flask, request, send_file, jsonify
from ultralytics import YOLO
from PIL import Image
import subprocess
import io
import glob
import os
app = Flask(__name__)
model = YOLO("runs/detect/train/weights/ best.pt ")
@app.route('/detect_faces', methods=['POST'])
def detect_faces():
if 'image' not in request.files:
return jsonify({"error": "No image provided"}), 400
image = Image.open (request.files ['image'] ).convert('RGB')
temp_image_path = "temp_image.jpg"
image.save(temp_image_path)
subprocess.run (f"yolo task=detect mode=predict model='runs/detect/train/weights/ best.pt ' conf=0.1 source={temp_image_path} save=True", shell=True)
predict_dirs = glob.glob("runs/detect/predict*/")
if predict_dirs:
latest_predict_dir = max(predict_dirs, key=os.path.getmtime)
img_path = os.path.join(latest_predict_dir, os.path.basename(temp_image_path))
img_path_with_new_extension = os.path.splitext(img_path)[0] + '.jpg'
result_image = Image.open (img_path_with_new_extension)
byte_io = io.BytesIO()
result_image.save(byte_io, 'JPEG')
byte_ io.seek (0)
return send_file(byte_io, mimetype='image/jpeg')
return jsonify({"error": "Prediction output not found."}), 500
@app.route('/', methods=['GET'])
def index():
return jsonify({"message": "API для обнаружения объектов. Отправьте POST запрос на /detect_faces с изображением."})
if __name__ == '__main__':
app.run (debug=True)
import sys
from PyQt5.QtWidgets import QApplication, QLabel, QPushButton, QVBoxLayout, QHBoxLayout, QFileDialog, QMainWindow,QWidget
from PyQt5.QtGui import QPixmap, QImage
import requests
from io import BytesIO
from PIL import Image
class FaceDetectionApp(QMainWindow):
def __init__(self):
super().__init__()
self.setWindowTitle("Распознавание лиц")
self.resize(800, 600)
self.init_ui()
def init_ui(self):
self.image_label = QLabel("Загрузите изображение")
self.image_label.setFixedSize(600, 400)
self.image_label.setStyleSheet("border: 1px solid black;")
self.image_label.setScaledContents(True)
self.upload_button = QPushButton("Загрузить изображение")
self.upload_button.clicked.connect(self.load_image)
self.detect_button = QPushButton("Отправить на сервер")
self.detect_button.clicked.connect(self.send_image)
self.detect_button.setEnabled(False)
layout = QVBoxLayout()
layout.addWidget(self.image_label)
layout.addWidget(self.upload_button)
layout.addWidget(self.detect_button)
container = QWidget()
container.setLayout(layout)
self.setCentralWidget(container)
def load_image(self):
options = QFileDialog.Options()
file_path, _ = QFileDialog.getOpenFileName(self, "Выберите изображение", "", "Images (*.png *.xpm *.jpg)",options=options)
if file_path:
self.image_path = file_path
pixmap = QPixmap(file_path)
self.image_label.setPixmap(pixmap)
self.detect_button.setEnabled(True)
def send_image(self):
url = "http://127.0.0.1:5000/detect_faces"
try:
with open(self.image_path, 'rb') as image_file:
response = requests.post(url, files={"image": image_file})
if response.status_code == 200:
content_type = response.headers.get('Content-Type')
if content_type == 'image/jpeg':
result_image = Image.open(BytesIO(response.content))
result_image = result_image.convert("RGBA")
data = result_image.tobytes("raw", "RGBA")
qimage = QImage(data, result_image.width, result_image.height, QImage.Format_RGBA8888)
pixmap = QPixmap.fromImage(qimage)
self.image_label.setPixmap(pixmap)
else:
self.image_label.setText("Ответ не является изображением")
else:
self.image_label.setText(f"Ошибка сервера: {response.status_code}")
except Exception as e:
self.image_label.setText(f"Ошибка: {str(e)}")
if __name__ == "__main__":
app = QApplication(sys.argv)
window = FaceDetectionApp()
window.show()
sys.exit(app.exec_())
много чего нужно