Compare commits

..

13 Commits

Author SHA1 Message Date
2e7cf69512 增加说明 2025-05-10 17:23:06 +08:00
76240a12e6 增加联邦学习评价指标。bugfix: 修复训练模型参数聚合问题 2025-05-10 17:22:56 +08:00
98321aa7d5 训练模型配置 2025-05-10 16:19:00 +08:00
d39aa31651 删除无用文件 2025-05-10 16:18:37 +08:00
f127ae2852 增加联邦学习指标;fix:Pytorch 加载模型不匹配 2025-05-07 10:41:36 +08:00
3a65d89315 ignore .vscode 2025-05-07 10:41:06 +08:00
2a3e5b17e7 yolov8对比训练 2025-05-05 17:30:12 +08:00
c57c8f3552 忽略训练结果和pt文件 2025-05-05 17:29:58 +08:00
310131d876 文件结构调整 2025-05-05 17:03:41 +08:00
myh
ba4508507b 评价指标优化 2025-04-22 21:41:58 +08:00
myh
89d8f4c0df 添加评价指标 2025-04-22 16:35:29 +08:00
myh
d1ed958db5 删除实例模块 2025-04-22 16:35:19 +08:00
myh
abd033b831 训练命令 2025-04-22 16:35:10 +08:00
11 changed files with 384 additions and 194 deletions

5
.gitignore vendored
View File

@@ -299,5 +299,8 @@ Temporary Items
# project files
/whl_packages/
/federated_learning/runs/detect/*
runs/
*.pt
*.cache
.vscode/
*.json

View File

@@ -1,3 +1,35 @@
# Graduation-Project
毕业设计基于YOLO和图像融合技术的无人机检测系统及安全性研究
毕业设计基于YOLO和图像融合技术的无人机检测系统及安全性研究
Linux 运行联邦训练
```bash
cd federated_learning
```
```bash
nohup python -u yolov8_fed.py > runtime.log 2>&1 &
```
Linux 运行集中训练
```bash
cd yolov8
```
```bash
nohup python -u yolov8_train.py > runtime.log 2>&1 &
```
实时监控日志文件
```bash
tail -f runtime.log
```
运行图像融合配准代码
```bash
cd image_fusion
```
```bash
python Image_Registration_test.py
```

View File

@@ -0,0 +1,49 @@
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
# Ultralytics YOLOv8 object detection model with P3/8 - P5/32 outputs
# Model docs: https://docs.ultralytics.com/models/yolov8
# Task docs: https://docs.ultralytics.com/tasks/detect
# Parameters
nc: 1 # number of classes
scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n'
# [depth, width, max_channels]
n: [0.33, 0.25, 1024] # YOLOv8n summary: 129 layers, 3157200 parameters, 3157184 gradients, 8.9 GFLOPS
s: [0.33, 0.50, 1024] # YOLOv8s summary: 129 layers, 11166560 parameters, 11166544 gradients, 28.8 GFLOPS
m: [0.67, 0.75, 768] # YOLOv8m summary: 169 layers, 25902640 parameters, 25902624 gradients, 79.3 GFLOPS
l: [1.00, 1.00, 512] # YOLOv8l summary: 209 layers, 43691520 parameters, 43691504 gradients, 165.7 GFLOPS
x: [1.00, 1.25, 512] # YOLOv8x summary: 209 layers, 68229648 parameters, 68229632 gradients, 258.5 GFLOPS
# YOLOv8.0n backbone
backbone:
# [from, repeats, module, args]
- [-1, 1, Conv, [64, 3, 2]] # 0-P1/2
- [-1, 1, Conv, [128, 3, 2]] # 1-P2/4
- [-1, 3, C2f, [128, True]]
- [-1, 1, Conv, [256, 3, 2]] # 3-P3/8
- [-1, 6, C2f, [256, True]]
- [-1, 1, Conv, [512, 3, 2]] # 5-P4/16
- [-1, 6, C2f, [512, True]]
- [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32
- [-1, 3, C2f, [1024, True]]
- [-1, 1, SPPF, [1024, 5]] # 9
# YOLOv8.0n head
head:
- [-1, 1, nn.Upsample, [None, 2, "nearest"]]
- [[-1, 6], 1, Concat, [1]] # cat backbone P4
- [-1, 3, C2f, [512]] # 12
- [-1, 1, nn.Upsample, [None, 2, "nearest"]]
- [[-1, 4], 1, Concat, [1]] # cat backbone P3
- [-1, 3, C2f, [256]] # 15 (P3/8-small)
- [-1, 1, Conv, [256, 3, 2]]
- [[-1, 12], 1, Concat, [1]] # cat head P4
- [-1, 3, C2f, [512]] # 18 (P4/16-medium)
- [-1, 1, Conv, [512, 3, 2]]
- [[-1, 9], 1, Concat, [1]] # cat head P5
- [-1, 3, C2f, [1024]] # 21 (P5/32-large)
- [[15, 18, 21], 1, Detect, [nc]] # Detect(P3, P4, P5)

View File

@@ -1,6 +1,9 @@
import glob
import os
from pathlib import Path
import json
from pydoc import cli
from threading import local
import yaml
from ultralytics import YOLO
@@ -15,121 +18,235 @@ def federated_avg(global_model, client_weights):
total_samples = sum(n for _, n in client_weights)
if total_samples == 0:
raise ValueError("Total number of samples must be positive.")
# DEBUG: global_dict
# print(global_model)
# 获取YOLO底层PyTorch模型参数
global_dict = global_model.model.state_dict()
# 提取所有客户端的 state_dict 和对应样本数
state_dicts, sample_counts = zip(*client_weights)
for key in global_dict:
# 对每一层参数取平均
# if global_dict[key].data.dtype == torch.float32:
# global_dict[key].data = torch.stack(
# [w[key].float() for w in client_weights], 0
# ).mean(0)
# 加权平均
if global_dict[key].dtype == torch.float32: # 只聚合浮点型参数
# 跳过 BatchNorm 层的统计量
if any(x in key for x in ['running_mean', 'running_var', 'num_batches_tracked']):
continue
# 按照样本数加权求和
weighted_tensors = [sd[key].float() * (n / total_samples)
for sd, n in zip(state_dicts, sample_counts)]
global_dict[key] = torch.stack(weighted_tensors, dim=0).sum(dim=0)
# 解决模型参数不匹配问题
try:
# 加载回YOLO模型
global_model.model.load_state_dict(global_dict)
except RuntimeError as e:
print('Ignoring "' + str(e) + '"')
# 添加调试输出
print("\n=== 参数聚合检查 ===")
# 选取一个典型参数层
# sample_key = list(global_dict.keys())[10]
# original = global_dict[sample_key].data.mean().item()
# aggregated = torch.stack([w[sample_key] for w in client_weights]).mean().item()
# print(f"参数层 '{sample_key}' 变化: {original:.4f} → {aggregated:.4f}")
# print(f"客户端参数差异: {[w[sample_key].mean().item() for w in client_weights]}")
# 随机选取一个非统计量层进行对比
sample_key = next(k for k in global_dict if 'running_' not in k)
aggregated_mean = global_dict[sample_key].mean().item()
client_means = [sd[sample_key].float().mean().item() for sd in state_dicts]
print(f"layer: '{sample_key}' Mean after aggregation: {aggregated_mean:.6f}")
print(f"The average value of the layer for each client: {client_means}")
# 克隆参数并脱离计算图
global_dict_copy = {
k: v.clone().detach().requires_grad_(False) for k, v in global_dict.items()
}
# 聚合可训练且存在的参数
for key in global_dict_copy:
# if global_dict_copy[key].dtype != torch.float32:
# continue
# if any(
# x in key for x in ["running_mean", "running_var", "num_batches_tracked"]
# ):
# continue
# 检查所有客户端是否包含当前键
all_clients_have_key = all(key in sd for sd in state_dicts)
if all_clients_have_key:
# 计算每个客户端的加权张量
# weighted_tensors = [
# client_state[key].float() * (sample_count / total_samples)
# for client_state, sample_count in zip(state_dicts, sample_counts)
# ]
weighted_tensors = []
for client_state, sample_count in zip(state_dicts, sample_counts):
weight = sample_count / total_samples # 计算权重
weighted_tensor = client_state[key].float() * weight # 加权张量
weighted_tensors.append(weighted_tensor)
# 聚合加权张量并更新全局参数
global_dict_copy[key] = torch.stack(weighted_tensors, dim=0).sum(dim=0)
# else:
# print(f"错误: 键 {key} 在部分客户端缺失,已保留全局参数")
# 终止训练或记录日志
# raise KeyError(f"键 {key} 缺失")
# 加载回YOLO模型
global_model.model.load_state_dict(global_dict_copy, strict=True)
# global_model.model.train()
# with torch.no_grad():
# global_model.model.load_state_dict(global_dict_copy, strict=True)
# 定义多个关键层
MONITOR_KEYS = [
"model.0.conv.weight",
"model.1.conv.weight",
"model.3.conv.weight",
"model.5.conv.weight",
"model.7.conv.weight",
"model.9.cv1.conv.weight",
"model.12.cv1.conv.weight",
"model.15.cv1.conv.weight",
"model.18.cv1.conv.weight",
"model.21.cv1.conv.weight",
"model.22.dfl.conv.weight",
]
with open("aggregation_check.txt", "a") as f:
f.write("\n=== 参数聚合检查 ===\n")
for key in MONITOR_KEYS:
# if key not in global_dict:
# continue
# if not all(key in sd for sd in state_dicts):
# continue
# 计算聚合后均值
aggregated_mean = global_dict[key].mean().item()
# 计算各客户端均值
client_means = [sd[key].float().mean().item() for sd in state_dicts]
with open("aggregation_check.txt", "a") as f:
f.write(f"'{key}' 聚合后均值: {aggregated_mean:.6f}\n")
f.write(f"各客户端该层均值差异: {[f'{cm:.6f}' for cm in client_means]}\n")
f.write(f"客户端最大差异: {max(client_means) - min(client_means):.6f}\n\n")
return global_model
# ------------ 修改训练流程 ------------
def federated_train(num_rounds, clients_data):
# ========== 初始化指标记录 ==========
metrics = {
"round": [],
"val_mAP": [], # 每轮验证集mAP
# "train_loss": [], # 每轮平均训练损失
"client_mAPs": [], # 各客户端本地模型在验证集上的mAP
"communication_cost": [], # 每轮通信开销MB
}
# 初始化全局模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
global_model = YOLO("yolov8n.pt").to(device)
# 设置类别数
global_model.model.nc = 1
global_model = (
YOLO("/home/image1325/DATA/Graduation-Project/federated_learning/yolov8n.yaml")
.load("/home/image1325/DATA/Graduation-Project/federated_learning/yolov8n.pt")
.to(device)
)
global_model.model.model[-1].nc = 1 # 设置检测类别数为1
# global_model.model.train.ema.enabled = False
# 克隆全局模型
local_model = copy.deepcopy(global_model)
for _ in range(num_rounds):
client_weights = []
# 各客户端的训练损失
# client_losses = []
# DEBUG: 检查全局模型参数
# global_dict = global_model.model.state_dict()
# print(global_dict.keys())
# 每个客户端本地训练
for data_path in clients_data:
# 统计本地训练样本数
with open(data_path, 'r') as f:
with open(data_path, "r") as f:
config = yaml.safe_load(f)
# Resolve img_dir relative to the YAML file's location
yaml_dir = os.path.dirname(data_path)
img_dir = os.path.join(yaml_dir, config.get('train', data_path)) # 从配置文件中获取图像目录
img_dir = os.path.join(
yaml_dir, config.get("train", data_path)
) # 从配置文件中获取图像目录
# print(f"Image directory: {img_dir}")
num_samples = (len(glob.glob(os.path.join(img_dir, '*.jpg'))) +
len(glob.glob(os.path.join(img_dir, '*.png'))))
num_samples = (
len(glob.glob(os.path.join(img_dir, "*.jpg")))
+ len(glob.glob(os.path.join(img_dir, "*.png")))
+ len(glob.glob(os.path.join(img_dir, "*.jpeg")))
)
# print(f"Number of images: {num_samples}")
# 克隆全局模型
local_model = copy.deepcopy(global_model)
local_model.model.load_state_dict(
global_model.model.state_dict(), strict=True
)
# 本地训练(保持你的原有参数设置)
local_model.train(
name=f"train{_ + 1}", # 当前轮次
data=data_path,
epochs=16, # 每轮本地训练1个epoch
save_period=16,
imgsz=640, # 图像大小
# model=local_model,
epochs=16, # 每轮本地训练多少个epoch
# save_period=16,
imgsz=768, # 图像大小
verbose=False, # 关闭冗余输出
batch=-1
batch=-1, # 批大小
workers=6, # 工作线程数
)
# 记录客户端训练损失
# client_loss = results.results_dict['train_loss']
# client_losses.append(client_loss)
# 收集模型参数及样本数
client_weights.append((copy.deepcopy(local_model.model.state_dict()), num_samples))
client_weights.append((local_model.model.state_dict(), num_samples))
# 聚合参数更新全局模型
global_model = federated_avg(global_model, client_weights)
print(f"Round {_ + 1}/{num_rounds} completed.")
return global_model
# DEBUG: 检查全局模型参数
# keys = global_model.model.state_dict().keys()
# ========== 评估全局模型 ==========
# 复制全局模型以避免在评估时修改参数
val_model = copy.deepcopy(global_model)
# 评估全局模型在验证集上的性能
with torch.no_grad():
val_results = val_model.val(
data="/mnt/DATA/uav_dataset_old/UAVdataset/fed_data.yaml", # 指定验证集配置文件
imgsz=768, # 图像大小
batch=16, # 批大小
verbose=False, # 关闭冗余输出
)
# 丢弃评估模型
del val_model
# DEBUG: 检查全局模型参数
# if keys != global_model.model.state_dict().keys():
# print("模型参数不一致!")
val_mAP = val_results.box.map # 获取mAP@0.5
# 计算平均训练损失
# avg_train_loss = sum(client_losses) / len(client_losses)
# 计算通信开销(假设传输全部模型参数)
model_size = sum(p.numel() * 4 for p in global_model.model.parameters()) / (
1024**2
) # MB
# 记录到指标容器
metrics["round"].append(_ + 1)
metrics["val_mAP"].append(val_mAP)
# metrics['train_loss'].append(avg_train_loss)
metrics["communication_cost"].append(model_size)
# 打印当前轮次结果
with open("aggregation_check.txt", "a") as f:
f.write(f"\n[Round {_ + 1}/{num_rounds}]\n")
f.write(f"Validation mAP@0.5: {val_mAP:.4f}\n")
# f.write(f"Average Train Loss: {avg_train_loss:.4f}")
f.write(f"Communication Cost: {model_size:.2f} MB\n\n")
return global_model, metrics
# ------------ 使用示例 ------------
if __name__ == "__main__":
# 联邦训练配置
clients_config = [
"/root/autodl-tmp/dataset/train1/train1.yaml", # 客户端1数据路径
"/root/autodl-tmp/dataset/train2/train2.yaml" # 客户端2数据路径
"/mnt/DATA/uav_fed/train1/train1.yaml", # 客户端1数据路径
"/mnt/DATA/uav_fed/train2/train2.yaml", # 客户端2数据路径
]
# 使用本地数据集进行测试
# clients_config = [
# "/home/image1325/DATA/Graduation-Project/dataset/train1/train1.yaml",
# "/home/image1325/DATA/Graduation-Project/dataset/train2/train2.yaml",
# ]
# 运行联邦训练
final_model = federated_train(num_rounds=10, clients_data=clients_config)
final_model, metrics = federated_train(num_rounds=10, clients_data=clients_config)
# 保存最终模型
final_model.save("yolov8n_federated.pt")
# final_model.export(format="onnx") # 导出为ONNX格式
# 检查1确认模型保存
# assert Path("yolov8n_federated.onnx").exists(), "模型导出失败"
# 检查2验证预测功能
# results = final_model.predict("../dataset/val/images/VS_P65.jpg", save=True)
# assert len(results[0].boxes) > 0, "预测结果异常"
with open("metrics.json", "w") as f:
json.dump(metrics, f, indent=4)

Binary file not shown.

View File

@@ -8,12 +8,41 @@ import cv2
import numpy as np
from ultralytics import YOLO
from skimage.metrics import structural_similarity as ssim
# 添加YOLOv8模型初始化
yolo_model = YOLO("yolov8n.pt") # 可替换为yolov8s/m/l等
yolo_model = YOLO("best.pt") # 可替换为yolov8s/m/l等
yolo_model.to('cuda') # 启用GPU加速
def calculate_en(img):
"""计算信息熵(处理灰度图)"""
hist = cv2.calcHist([img], [0], None, [256], [0, 256])
hist = hist / hist.sum()
return -np.sum(hist * np.log2(hist + 1e-10))
def calculate_sf(img):
"""计算空间频率(处理灰度图)"""
rf = np.sqrt(np.mean(np.square(np.diff(img, axis=0))))
cf = np.sqrt(np.mean(np.square(np.diff(img, axis=1))))
return np.sqrt(rf ** 2 + cf ** 2)
def calculate_mi(img1, img2):
"""计算互信息(处理灰度图)"""
hist_2d = np.histogram2d(img1.ravel(), img2.ravel(), 256)[0]
pxy = hist_2d / hist_2d.sum()
px = np.sum(pxy, axis=1)
py = np.sum(pxy, axis=0)
return np.sum(pxy * np.log2(pxy / (px[:, None] * py[None, :] + 1e-10) + 1e-10))
def calculate_ssim(img1, img2):
"""计算SSIM处理灰度图"""
return ssim(img1, img2, data_range=255)
# 裁剪线性RGB对比度拉伸去掉2%百分位以下的数去掉98%百分位以上的数,上下百分位数一般相同,并设置输出上下限)
def truncated_linear_stretch(image, truncated_value=2, maxout=255, min_out=0):
"""
@@ -145,32 +174,60 @@ def main(matchimg_vi, matchimg_in):
orimg_vi = matchimg_vi
orimg_in = matchimg_in
h, w = orimg_vi.shape[:2] # 480 640
flag, H, dot = Images_matching(matchimg_vi, matchimg_in) # (3, 3)//获取对应的配准坐标点
# (3, 3)//获取对应的配准坐标点
flag, H, dot = Images_matching(matchimg_vi, matchimg_in)
if flag == 0:
return 0, None, 0
return 0, None, 0, 0.0, 0.0, 0.0, 0.0
else:
# 配准处理
matched_ni = cv2.warpPerspective(orimg_in, H, (w, h))
matched_ni, left, right, top, bottom = removeBlackBorder(matched_ni)
# 裁剪可见光图像
# fusion = fusions(orimg_vi[left:right, top:bottom], matched_ni)
# 不裁剪可见光图像
fusion = fusions(orimg_vi, matched_ni)
# 转换为灰度计算指标
fusion_gray = cv2.cvtColor(fusion, cv2.COLOR_RGB2GRAY)
cropped_vi_gray = cv2.cvtColor(orimg_vi, cv2.COLOR_BGR2GRAY)
matched_ni_gray = matched_ni # 红外图已经是灰度
# 计算指标
en = calculate_en(fusion_gray)
sf = calculate_sf(fusion_gray)
mi_visible = calculate_mi(fusion_gray, cropped_vi_gray)
mi_infrared = calculate_mi(fusion_gray, matched_ni_gray)
mi_total = mi_visible + mi_infrared
# 添加SSIM容错处理
try:
ssim_visible = calculate_ssim(fusion_gray, cropped_vi_gray)
ssim_infrared = calculate_ssim(fusion_gray, matched_ni_gray)
ssim_avg = (ssim_visible + ssim_infrared) / 2
except Exception as ssim_error:
print(f"SSIM计算错误: {ssim_error}")
ssim_avg = -1 # 用-1表示计算失败
# YOLOv8目标检测
results = yolo_model(fusion) # 输入融合后的图像
annotated_image = results[0].plot() # 绘制检测框
return 1, annotated_image, dot # 返回带检测结果的图像
# 返回带检测结果的图像
return 1, annotated_image, dot, en, sf, mi_total, ssim_avg
except Exception as e:
print(f"Error in fusion/detection: {e}")
return 0, None, 0
return 0, None, 0, 0.0, 0.0, 0.0, 0.0
def parse_args():
# 输入可见光和红外图像路径
visible_image_path = "test/visible.jpg" # 可见光图片路径
infrared_image_path = "test/infrared.jpg" # 红外图片路径
visible_image_path = "./test/visible/visibleI0195.jpg" # 可见光图片路径
infrared_image_path = "./test/infrared/infraredI0195.jpg" # 红外图片路径
# 输入可见光和红外视频路径
visible_video_path = "test/visible.mp4" # 可见光视频路径
infrared_video_path = "test/infrared.mp4" # 红外视频路径
visible_video_path = "./test/visible.mp4" # 可见光视频路径
infrared_video_path = "./test/infrared.mp4" # 红外视频路径
"""解析命令行参数"""
parser = argparse.ArgumentParser(description='图像融合与目标检测')
@@ -272,13 +329,26 @@ if __name__ == '__main__':
img_inf_gray = cv2.cvtColor(img_infrared, cv2.COLOR_BGR2GRAY)
# 执行融合与检测
flag, fusion_result, _ = main(img_visible, img_inf_gray)
flag, fusion_result, dot, en, sf, mi, ssim_val = main(img_visible, img_inf_gray)
if flag == 1:
# 展示评价指标
print("\n======== 融合质量评价 ========")
print(f"信息熵EN: {en:.2f}")
print(f"空间频率SF: {sf:.2f}")
print(f"互信息MI: {mi:.2f}")
# 条件显示SSIM
if ssim_val >= 0:
print(f"结构相似性SSIM: {ssim_val:.4f}")
else:
print("结构相似性SSIM: 计算失败(已跳过)")
print(f"配准点数: {dot}")
# 显示并保存结果
cv2.imshow("Fusion with Detection", fusion_result)
# cv2.imshow("Fusion with Detection", fusion_result)
cv2.imwrite("output/fusion_result.jpg", fusion_result)
cv2.waitKey(0)
cv2.destroyAllWindows()
# cv2.waitKey(0)
# cv2.destroyAllWindows()
else:
print("融合失败!")

View File

@@ -1,74 +0,0 @@
import numpy as np
import cv2
from skimage.metrics import structural_similarity as ssim
from skimage.filters import sobel
from sklearn.metrics import mutual_info_score
# Helper to compute mutual information between two grayscale images
def evaluate_mutual_information(img1_gray, img2_gray):
hist_2d, _, _ = np.histogram2d(img1_gray.ravel(), img2_gray.ravel(), bins=256)
pxy = hist_2d / float(np.sum(hist_2d))
px = np.sum(pxy, axis=1)
py = np.sum(pxy, axis=0)
px_py = np.outer(px, py)
nzs = pxy > 0
mi = np.sum(pxy[nzs] * np.log(pxy[nzs] / px_py[nzs]))
return mi
# Compute SSIM between two grayscale images
def evaluate_registration_ssim(img1_gray, img2_gray):
return ssim(img1_gray, img2_gray)
# Entropy of grayscale image (fusion quality)
def evaluate_fusion_entropy(fusion_img):
gray = cv2.cvtColor(fusion_img, cv2.COLOR_RGB2GRAY)
hist = cv2.calcHist([gray], [0], None, [256], [0, 256])
hist = hist.ravel() / hist.sum()
entropy = -np.sum(hist * np.log2(hist + 1e-9))
return entropy
# Edge strength using Sobel (fusion quality)
def evaluate_fusion_edges(fusion_img):
gray = cv2.cvtColor(fusion_img, cv2.COLOR_RGB2GRAY)
edges = sobel(gray.astype(float) / 255.0)
return np.mean(edges)
# SSIM between fused image and one of the sources
def evaluate_fusion_ssim(fusion_img, reference_img):
fusion_gray = cv2.cvtColor(fusion_img, cv2.COLOR_RGB2GRAY)
ref_gray = cv2.cvtColor(reference_img, cv2.COLOR_RGB2GRAY)
return ssim(fusion_gray, ref_gray)
# Return all in one place (stub images would be required to test)
def summarize_evaluation(img1_gray, img2_gray, fusion_img, ref_img_for_ssim):
return {
"Registration SSIM": evaluate_registration_ssim(img1_gray, img2_gray),
"Mutual Information": evaluate_mutual_information(img1_gray, img2_gray),
"Fusion Entropy": evaluate_fusion_entropy(fusion_img),
"Fusion Edge Strength": evaluate_fusion_edges(fusion_img),
"Fusion SSIM (vs Ref)": evaluate_fusion_ssim(fusion_img, ref_img_for_ssim),
}
# 将所有评价封装成一个高层函数 evaluate_all
def evaluate_all(img1_gray, img2_gray, fusion_img, ref_img_for_ssim, verbose=True):
"""
评估图像配准和融合质量的通用函数
:param img1_gray: 可见光灰度图像(原图)
:param img2_gray: 红外灰度图像(配准后)
:param fusion_img: 融合图像RGB
:param ref_img_for_ssim: 可见光RGB图用于对比SSIM
:param verbose: 是否打印结果
:return: dict 评价指标结果
"""
results = summarize_evaluation(img1_gray, img2_gray, fusion_img, ref_img_for_ssim)
if verbose:
print("图像评价指标如下:")
for k, v in results.items():
print(f"{k}: {v:.4f}")
return results

View File

@@ -1,26 +0,0 @@
from evaluate import *
# 创建模拟图像数据用于测试
# img1_gray原始灰度图像可见光
# img2_gray变换后的灰度图像红外模拟
# fusion_img融合图像可见光 + 红外)
# ref_img_for_ssim参考图像可见光RGB
# 创建基础灰度图像(梯度)
img1_gray = np.tile(np.linspace(50, 200, 256).astype(np.uint8), (256, 1))
# 模拟配准后的图像:加一点噪声和平移
img2_gray = np.roll(img1_gray, shift=5, axis=1) # 平移模拟配准偏差
noise = np.random.normal(0, 5, img2_gray.shape).astype(np.uint8)
img2_gray = cv2.add(img2_gray, noise)
# 创建 RGB 可见光图(重复三个通道)
ref_img_for_ssim = cv2.merge([img1_gray] * 3)
# 创建融合图像取两个灰度图平均后合并入RGB
fusion_Y = cv2.addWeighted(img1_gray, 0.5, img2_gray, 0.5, 0)
fusion_img = cv2.merge([fusion_Y, img1_gray, img2_gray])
# 运行评价函数
scores = evaluate_all(img1_gray, img2_gray, fusion_img, ref_img_for_ssim)

Binary file not shown.

6
yolov8/yolov8.yaml Normal file
View File

@@ -0,0 +1,6 @@
train: /mnt/DATA/dataset/uav_dataset/train/images/
val: /mnt/DATA/dataset/uav_dataset/val/images/
test: /mnt/DATA/dataset/test2/images/
# number of classes
nc: 1
names: ['uav']

13
yolov8/yolov8_train.py Normal file
View File

@@ -0,0 +1,13 @@
from ultralytics import YOLO
# 加载预训练模型
model = YOLO('../yolov8n.pt')
# 开始训练
model.train(
data='./yolov8.yaml', # 数据配置文件路径
epochs=320, # 训练轮数
batch=-1, # 批量大小
imgsz=640, # 输入图片大小
device=0 # 使用的设备0 表示 GPU'cpu' 表示 CPU
)