jiangchaoqing
/
tuoheng_AIPlatform
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
tuoheng_AIPlatform/yolov5-th/yolov5_kafka.py

463 lines
13 KiB
Python
Raw Permalink Normal View History

code
2025-07-15 10:01:04 +08:00
# -*- coding: utf-8 -*-
"""
HTTP to Kafka proxy service for training requests with message transformation
Features:
- Receives HTTP POST requests with training parameters
- Transforms message format from incoming to Kafka-compatible format
- Validates the request format
- Forwards valid requests to Kafka
- Returns appropriate HTTP responses
"""
from flask import Flask, request, jsonify
import os
from kafka import KafkaProducer
import json
import random
import string
import logging
from multiprocessing import Process
import time
import os
import random
import yaml
from shutil import copyfile
from minio import Minio
from minio.error import S3Error
from minio import Minio
from minio.error import S3Error
import os
import urllib.parse
from minio import Minio
from minio.error import S3Error
import os
import urllib.parse
app = Flask(__name__)
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Kafka configuration
KAFKA_CONFIG = {
'bootstrap_servers': '172.18.0.126:9094',
'value_serializer': lambda v: json.dumps(v).encode('utf-8')
}
# Kafka topics
TOPIC_TRAIN = 'training-tasks'
model_dict = {
"001": "AnglerSwimmer",
"002": "channel2",
"003": "channelEmergency",
"004": "cityMangement",
"005": "cityMangement2",
"006": "cityMangement3",
"007": "cityRoad",
"008": "conf",
"009": "countryRoad",
"010": "crackMeasurement",
"011": "crowdCounting",
"012": "drowning",
"013": "forest",
"014": "forest2",
"015": "forestCrowd",
"016": "highWay2",
"017": "illParking",
"018": "noParking",
"019": "ocr2",
"020": "ocr_en",
"021": "pedestrian",
"022": "pothole",
"023": "river",
"024": "river2",
"025": "riverT",
"026": "road",
"027": "ship",
"028": "ship2",
"029": "smogfire",
"030": "trafficAccident",
"031": "vehicle"
}
def validate_minio_url(url):
"""验证MinIO URL格式是否正确"""
try:
parsed = urllib.parse.urlparse(url)
if not all([parsed.scheme in ('http', 'https'), parsed.netloc, len(parsed.path.split('/')) >= 3]):
raise ValueError("Invalid MinIO URL format")
return parsed
except Exception as e:
raise ValueError(f"URL validation failed: {str(e)}")
def download_minio_directory(minio_url, local_base_dir):
"""
下载MinIO目录到本地
:param minio_url: 完整的MinIO路径 (e.g. "http://172.18.0.129:8084/th-dsp/027123")
:param local_base_dir: 本地存储基础目录 (e.g. "/home/th/jcq/AI_AutoPlat/DATA/TrainDatasets")
:return: (success, message) 元组
"""
try:
# 验证并解析URL
parsed = validate_minio_url(minio_url)
endpoint = parsed.netloc
bucket_name = parsed.path.split('/')[1]
prefix = '/'.join(parsed.path.split('/')[2:])
# 从环境变量获取凭证（推荐方式）
access_key = os.getenv("MINIO_ACCESS_KEY", "IKf3A0ZSXsR1m0oalMjV") # 默认值仅用于测试
secret_key = os.getenv("MINIO_SECRET_KEY", "yoC6qRo2hlyZu8Pdbt6eh9TVaTV4gD7KRudromrk")
# 创建本地目录（使用prefix的最后部分作为目录名）
dir_name = os.path.basename(prefix) or "minio_download"
local_dir = os.path.join(local_base_dir, dir_name)
os.makedirs(local_dir, exist_ok=True)
# 初始化MinIO客户端
client = Minio(
endpoint,
access_key=access_key,
secret_key=secret_key,
secure=parsed.scheme == 'https'
)
print(f"Starting download from {minio_url} to {local_dir}")
# 递归下载所有文件
objects = client.list_objects(bucket_name, prefix=prefix, recursive=True)
downloaded_files = 0
for obj in objects:
# 构建本地路径（保持相对目录结构）
relative_path = os.path.relpath(obj.object_name, prefix)
local_path = os.path.join(local_dir, relative_path)
# 确保目录存在
os.makedirs(os.path.dirname(local_path), exist_ok=True)
# 下载文件
client.fget_object(bucket_name, obj.object_name, local_path)
downloaded_files += 1
print(f"Downloaded [{downloaded_files}]: {obj.object_name} -> {local_path}")
print(f"Download completed! Total files: {downloaded_files}")
return True, local_dir
except S3Error as e:
error_msg = f"MinIO Error: {str(e)}"
if "NoSuchBucket" in str(e):
error_msg = f"Bucket not found: {bucket_name}"
return False, error_msg
except ValueError as e:
return False, f"Invalid URL: {str(e)}"
except Exception as e:
return False, f"Unexpected error: {str(e)}"
def split_dataset(data_dir, train_ratio=0.8):
# 获取所有jpg文件
jpg_files = [f for f in os.listdir(data_dir) if f.endswith('.jpg')]
# 检查对应的txt文件是否存在
valid_files = []
for jpg in jpg_files:
txt = jpg.replace('.jpg', '.txt')
if os.path.exists(os.path.join(data_dir, txt)):
valid_files.append(jpg)
# 随机打乱
random.shuffle(valid_files)
# 划分训练集和验证集
split_idx = int(len(valid_files) * train_ratio)
train_files = valid_files[:split_idx]
val_files = valid_files[split_idx:]
return train_files, val_files
def write_file_list(file_list, output_file, base_dir):
with open(output_file, 'w') as f:
for file in file_list:
f.write(os.path.join(base_dir, file) + '\n')
def copy_files_to_folders(files, src_dir, dest_dir):
"""复制图片和标签文件到目标文件夹"""
os.makedirs(dest_dir, exist_ok=True)
for file in files:
# 复制图片
src_img = os.path.join(src_dir, file)
dest_img = os.path.join(dest_dir, file)
copyfile(src_img, dest_img)
# 复制对应的标签文件
src_txt = os.path.join(src_dir, file.replace('.jpg', '.txt'))
dest_txt = os.path.join(dest_dir, file.replace('.jpg', '.txt'))
copyfile(src_txt, dest_txt)
def update_yaml(yaml_file, train_txt, val_txt, nc, names):
with open(yaml_file, 'r') as f:
data = yaml.safe_load(f)
data['train'] = train_txt
data['val'] = val_txt
data['nc'] = nc
data['names'] = names
with open(yaml_file, 'w') as f:
yaml.dump(data, f, default_flow_style=False)
def get_model_name(code):
# Ensure code is a 3-digit string
code_str = f"{int(code):03d}" if str(code).isdigit() else code
return model_dict.get(code_str, f"Unknown code: {code}")
def generate_request_id(length=30):
"""Generate a unique request ID"""
return 'bb' + ''.join(random.sample(string.ascii_letters, length))
def transform_message_format(input_data):
transformed = {
"Scene": input_data.get("ModelScenetype", "train").capitalize(),
"Command": "start", # Default command
"Request_ID": input_data.get("request_id", generate_request_id()),
"Version": input_data.get("Version", "v1.0"),
"Name": "channel2", # Default name if not provided
"Model": input_data.get("code", input_data.get("Code", "026")),
}
# Handle TrainParameter construction
params = input_data.get("parameters", {})
batch_size = params.get("batch_size")
img_size = params.get("img_size")
epochs = params.get("epochs")
model_name = get_model_name(input_data.get("code", input_data.get("Code", "026")))
RAW_YAML_PATH = "/home/th/jcq/AI_AutoPlat/AI_web_dsj/config/yaml"
yaml_path = os.path.join(RAW_YAML_PATH,f"{model_name}.yaml")
print("####line137",yaml_path)
# minio_url = "http://172.18.0.129:8084/th-dsp/027123"
# 获取URL
minio_url = input_data.get("TrainImageDir")
if not minio_url:
print("Error: TrainImageDir not specified in input_data")
else:
# 指定本地存储目录
local_base_dir = "/home/th/jcq/AI_AutoPlat/DATA/TrainDatasets"
# 执行下载
success, result = download_minio_directory(minio_url, local_base_dir)
if success:
print(f"Successfully downloaded to: {result}")
else:
print(f"Download failed: {result}")
if success:
print("line277",f"Files downloaded successfully to: {result}")
else:
print(f"Download failed: {result}")
TrainImageDir = result
# 创建 train 和 val 文件夹
train_dir = os.path.join(TrainImageDir, "train")
val_dir = os.path.join(TrainImageDir, "val")
os.makedirs(train_dir, exist_ok=True)
os.makedirs(val_dir, exist_ok=True)
# 划分数据集
train_files, val_files = split_dataset(TrainImageDir)
# 复制文件到 train/ 和 val/ 文件夹
copy_files_to_folders(train_files, TrainImageDir, train_dir)
copy_files_to_folders(val_files, TrainImageDir, val_dir)
# 写入训练集和验证集txt文件
train_txt = os.path.join(TrainImageDir, "train.txt")
val_txt = os.path.join(TrainImageDir, "val.txt")
write_file_list(train_files, train_txt, train_dir)
write_file_list(val_files, val_txt, val_dir)
# 更新yaml文件
nc = 5
names = ['flag', 'buoy', 'shipname', 'ship', 'uncover']
update_yaml(yaml_path, train_txt, val_txt, nc, names)
print(f"数据集划分完成，训练集: {len(train_files)} 张，验证集: {len(val_files)} 张")
print(f"训练集列表已写入: {train_txt}")
print(f"验证集列表已写入: {val_txt}")
print(f"YAML文件已更新: {yaml_path}")
RAW_WEIGHTS_PATH = "/home/th/jcq/AIlib/weights_init"
model_path = os.path.join(RAW_WEIGHTS_PATH,model_name,"yolov5.pt")
print("####line98",model_path)
transformed["TrainParameter"] = str([
batch_size,
img_size,
epochs,
yaml_path,
model_path
])
return transformed
def validate_training_request(data):
"""Validate the training request structure"""
required_fields = {
"Scene": str,
"Command": str,
"Version": str,
"Name": str,
"Model": str,
"TrainParameter": str
}
for field, field_type in required_fields.items():
if field not in data:
return False, f"Missing required field: {field}"
if not isinstance(data[field], field_type):
return False, f"Field {field} must be {field_type.__name__}"
if data["Scene"] not in ["Train"]:
return False, "Invalid Scene value"
if data["Command"] not in ["start", "stop"]:
return False, "Invalid Command value"
return True, ""
def send_to_kafka(topic, message):
"""Send message to Kafka topic"""
producer = None
try:
producer = KafkaProducer(**KAFKA_CONFIG)
future = producer.send(topic, value=message)
# Wait for message to be delivered
result = future.get(timeout=10)
logger.info(f"Message sent to partition {result.partition}, offset {result.offset}")
return True, ""
except Exception as e:
logger.error(f"Failed to send message: {e}")
return False, str(e)
finally:
if producer:
producer.flush()
producer.close()
@app.route('/api/train', methods=['POST'])
def handle_training_request():
"""HTTP endpoint for training requests (non-blocking)"""
try:
# Parse and validate basic request structure
data = request.get_json()
if not data or not isinstance(data, dict):
return jsonify({"status": "error", "message": "Invalid JSON data"}), 400
# Generate request ID
request_id = generate_request_id()
# Start background process (daemon=True ensures process exits with main)
p = Process(
target=process_training_task,
args=(data, request_id),
daemon=True
)
p.start()
# Immediate response
return jsonify({
"status": "success",
"message": "Training task started",
"request_id": request_id,
"process_id": p.pid
}), 202 # 202 Accepted
except Exception as e:
logger.error(f"Request handling error: {str(e)}")
return jsonify({
"status": "error",
"message": "Internal server error"
}), 500
def process_training_task(data, request_id):
"""独立进程执行耗时任务"""
try:
logger.info(f"[{request_id}] Starting background task (PID: {os.getpid()})")
start_time = time.time()
# 1. 数据转换
transformed = transform_message_format(data)
transformed["Request_ID"] = request_id
# 2. 数据验证
is_valid, msg = validate_training_request(transformed)
if not is_valid:
raise ValueError(msg)
# 3. 发送到Kafka
success, error = send_to_kafka(TOPIC_TRAIN, transformed)
if not success:
raise RuntimeError(error)
elapsed = time.time() - start_time
logger.info(f"[{request_id}] Task completed in {elapsed:.2f}s")
except Exception as e:
logger.error(f"[{request_id}] Task failed: {str(e)}")
if __name__ == '__main__':
app.run(host='192.168.10.11', port=6688, debug=True)