# 导入必要的库
import tensorflow as tf
from tensorflow.keras import layers, models, optimizers
import numpy as np

# 定义加载模型并替换全连接层的函数
def load_model(model_name, num_classes):
    """
    加载预训练模型（VGG16或ResNet50），冻结卷积层，并替换顶层的全连接层。
    
    参数：
        model_name (str): 要使用的模型名称，可选 'VGG16' 或 'ResNet50'
        num_classes (int): 分类任务的目标类别数
    
    返回：
        model (tf.keras.Model): 修改后的Keras模型
    """
    if model_name == 'VGG16':
        # 使用ImageNet预训练权重初始化VGG16模型
        # include_top=False 表示不包含原始的全连接层
        # input_shape=(224, 224, 3) 设置输入图像尺寸为224x224的RGB图像
        # pooling='avg' 表示使用全局平均池化替代最后的空间维度
        base_model = tf.keras.applications.VGG16(
            weights='imagenet',
            include_top=False,
            input_shape=(224, 224, 3),
            pooling='avg'
        )

    elif model_name == 'ResNet50':
        # 使用ImageNet预训练权重初始化ResNet50模型
        base_model = tf.keras.applications.ResNet50(
            weights='imagenet',
            include_top=False,
            input_shape=(224, 224, 3),
            pooling='avg'
        )
    
    # 冻结base_model中的所有层，即在训练时不更新这些参数
    base_model.trainable = False

    # 构建新的模型：将base_model作为基础，在其顶部添加新的全连接层
    model = models.Sequential([
        base_model,  # 基础模型部分（卷积层）
        layers.Dense(num_classes, activation='softmax')  # 新的输出层，用于分类
    ])
    
    # 打印模型结构摘要
    model.summary()
    
    return model


# 超参数设置
num_classes = 3               # 分类数量，例如三分类问题
model_name = 'VGG16'          # 使用的模型名称，也可以是'ResNet50'
learning_rate = 0.001         # 学习率
batch_size = 32               # 每个批次的数据量
image_size = (224, 224)       # 图像输入尺寸，适配大多数预训练模型

# 加载并构建模型
model = load_model(model_name, num_classes)

# 接下来可以继续编译和训练模型，例如：
# model.compile(optimizer=optimizers.Adam(learning_rate=learning_rate),
#               loss='categorical_crossentropy',
#               metrics=['accuracy'])

# 然后准备数据集并调用 model.fit(...) 进行训练

import torch
import torch.nn as nn
import torchvision.models as models
from torchvision.models import VGG16_Weights, ResNet18_Weights


def load_model(model_name, num_classes):
    """
    加载指定的预训练模型，并修改最后的分类层以适应新的分类任务。

    参数:
        model_name (str): 模型名称 ('vgg16' 或 'resnet18')
        num_classes (int): 新的分类任务的类别数量

    返回:
        model (torch.nn.Module): 修改后的模型
    """

    if model_name == 'vgg16':
        # 加载预训练的 VGG16 模型
        model = models.vgg16(weights=VGG16_Weights.DEFAULT)

        # 冻结特征提取层（不计算梯度）
        for param in model.features.parameters():
            param.requires_grad = False

        # 通过一次前向传播获取特征输出维度
        with torch.no_grad():
            dummy_input = torch.randn(1, 3, 224, 224)
            features = model.features(dummy_input)
            features_dim = features.view(features.size(0), -1).size(1)  # 展平后的维度

        # 替换最后的全连接层为自定义的分类器
        model.classifier = nn.Sequential(
            nn.Linear(features_dim, 256),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(256, num_classes)
        )

    elif model_name == 'resnet18':
        # 加载预训练的 ResNet18 模型
        model = models.resnet18(weights=ResNet18_Weights.DEFAULT)

        # 冻结卷积层参数（只训练最后的全连接层）
        for name, param in model.named_parameters():
            if not isinstance(model._modules.get(name), nn.Linear):
                param.requires_grad = False

        # 获取最后一个全连接层的输入维度
        in_features = model.fc.in_features

        # 替换最后的全连接层为自定义的分类器
        model.fc = nn.Sequential(
            nn.Linear(in_features, 256),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(256, num_classes)
        )

    else:
        raise ValueError(f'Unknown model name: {model_name}')

    # 打印模型结构
    print(model)

    return model

# 超参数
num_classes = 3
model_name = 'vgg16'
model = load_model(model_name, num_classes)

import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np


def load_image(path):
    """
    加载图像文件并将其转换为TensorFlow张量

    参数:
        path: 字符串，图像文件路径

    返回:
        归一化到[0,1]范围的float32类型图像张量
    """
    # 读取图像文件
    image = tf.io.read_file(path)
    # 解码JPEG图像(3通道RGB)
    decode_image = tf.image.decode_jpeg(image, channels=3)
    # 将图像转换为float32类型并归一化到[0,1]
    img = tf.image.convert_image_dtype(decode_image, tf.float32)
    return img


def display_images(image_list=None, title=None):
    """
    显示图像网格

    参数:
        image_list: 图像列表，支持多种嵌套格式
        title: 图像显示的标题
    """
    if image_list is None:
        image_list = []

    if title is None:
        title = ""

    # 检查是否是嵌套列表，如果不是则转换为嵌套列表
    if not any(isinstance(i, list) for i in image_list):
        image_list = [image_list]  # 例如 [img, img, img] => [[img, img, img]]

    # 计算行列数
    rows = len(image_list)
    cols = max(len(row) if isinstance(row, list) else 1 for row in image_list)

    # 创建子图
    plt.suptitle(title)
    fig, ax = plt.subplots(rows, cols)

    # 确保ax是2D数组(统一处理单行/单列情况)
    ax = np.atleast_2d(ax)

    # 遍历每个图像并显示
    for i, row in enumerate(image_list):
        if not isinstance(row, list):
            row = [row]

        for j, img in enumerate(row):
            ax[i, j].imshow(img)
            ax[i, j].axis("off")  # 关闭坐标轴

        # 填充空白子图
        for j in range(len(row), cols):
            ax[i, j].axis("off")

    plt.tight_layout()  # 自动调整子图间距
    plt.show()


def resize_center_crop(img, target_size=224):
    """
    保持宽高比调整大小后中心裁剪

    参数:
        img: 输入图像张量
        target_size: 目标尺寸(正方形)

    返回:
        中心裁剪后的图像
    """
    height, width, ch = img.shape
    aspect_ratio = width / height

    # 计算保持宽高比的临时尺寸
    if aspect_ratio > 1:  # 宽>高
        target_height = target_size
        target_width = int(round(target_size * aspect_ratio))
    else:  # 高>=宽
        target_width = target_size
        target_height = int(round(target_size / aspect_ratio))

    # 调整图像大小
    scale_img = tf.image.resize(img, [target_height, target_width])

    # 计算裁剪偏移量
    offset_height = (target_height - target_size) // 2
    offset_width = (target_width - target_size) // 2

    # 中心裁剪
    crop_img = tf.image.crop_to_bounding_box(
        scale_img,
        offset_height,
        offset_width,
        target_size,
        target_size
    )

    return crop_img


def resize_pad_crop(img, target_size=224):
    """
    保持宽高比调整大小后填充为正方形

    参数:
        img: 输入图像张量
        target_size: 目标尺寸(正方形)

    返回:
        填充后的图像
    """
    height, width, ch = img.shape
    aspect_ratio = width / height

    # 计算保持宽高比的临时尺寸
    if aspect_ratio > 1:  # 宽>高
        target_width = target_size
        target_height = int(round(target_size / aspect_ratio))
    else:  # 高>=宽
        target_height = target_size
        target_width = int(round(target_size * aspect_ratio))

    # 调整图像大小
    scale_img = tf.image.resize(img, [target_height, target_width])

    # 计算需要填充的像素数
    pad_height = target_size - target_height
    pad_width = target_size - target_width

    # 计算上下左右的填充量
    pad_top = pad_height // 2
    pad_bottom = pad_height - pad_top
    pad_left = pad_width // 2
    pad_right = pad_width - pad_left

    print(f"填充量: 上{pad_top}, 下{pad_bottom}, 左{pad_left}, 右{pad_right}")
    # 定义填充模式(只在高度和宽度维度填充)
    padding = [[pad_top, pad_bottom], [pad_left, pad_right], [0, 0]]
    # 执行填充(用0填充)
    pad_img = tf.pad(scale_img, padding, "CONSTANT", constant_values=0)

    return pad_img


def gaussian_noise(img, mean=0, stddev=0.1):
    """
    添加高斯噪声

    参数:
        img: 输入图像
        mean: 噪声均值
        stddev: 噪声标准差

    返回:
        添加噪声后的图像
    """
    # 生成高斯噪声
    noise = tf.random.normal(
        shape=tf.shape(img),
        mean=mean,
        stddev=stddev,
        dtype=tf.float32
    )
    # 添加噪声并裁剪到[0,1]范围
    noise_img = tf.clip_by_value(img + noise, 0, 1)
    return noise_img


def salt_and_pepper_noise(img, prob=0.1):
    """
    添加椒盐噪声

    参数:
        img: 输入图像
        prob: 噪声像素比例

    返回:
        添加噪声后的图像
    """
    height, width, chs = img.shape
    # 计算噪声像素数量
    num_noise_pixels = int(height * width * prob)

    # 随机生成噪声像素坐标
    y_coords = np.random.randint(0, height, num_noise_pixels)
    x_coords = np.random.randint(0, width, num_noise_pixels)

    # 随机决定是盐噪声(1)还是胡椒噪声(0)
    salt_or_pepper = np.random.choice([0, 1], size=num_noise_pixels * chs)
    salt_or_pepper = salt_or_pepper.reshape(num_noise_pixels, chs)

    # 创建图像副本并添加噪声
    noise_img = img.numpy().copy()
    for i in range(num_noise_pixels):
        noise_img[y_coords[i], x_coords[i]] = salt_or_pepper[i]

    return noise_img


def random_mask(img, mask_size=50, num_blocks=1):
    """
    随机添加黑色遮挡块

    参数:
        img: 输入图像
        mask_size: 遮挡块大小
        num_blocks: 遮挡块数量

    返回:
        添加遮挡后的图像
    """
    height, width, chs = img.shape
    img_np = img.numpy()
    # 添加指定数量的遮挡块
    for _ in range(num_blocks):
        # 随机生成遮挡块左上角坐标
        y = np.random.randint(0, height - mask_size)
        x = np.random.randint(0, width - mask_size)
        # 将区域设置为黑色(0)
        img_np[y:y + mask_size, x:x + mask_size] = 0

    return img_np

img = load_image("./kpbl.PNG")
#水平翻转
img_flipped = tf.image.flip_left_right(img)
#上下翻转
img_updown = tf.image.flip_up_down(img)
#旋转
img_rotated = tf.image.rot90(img, k=4)
#缩放
img_resized = tf.image.resize(img, size=[300, 300])
#按中心裁剪
img_cropped = tf.image.central_crop(img, central_fraction=0.5)
#按尺寸进行裁剪
img_cropped = tf.image.crop_to_bounding_box(img, offset_height=100, offset_width=100, target_height=200, target_width=200)
#随机裁剪，使用随机种子
img_cropped = tf.image.random_crop(img, size=[400, 400, 3], seed=1)

# 颜色相关的增强
# 增加亮度
img_bright = tf.image.adjust_brightness(img, delta=0.3) #delta > 0，它会增加亮度，delta < 0，它会减少亮度
# 增加对比度
img_contrast1 = tf.image.adjust_contrast(img, contrast_factor=2)# contrast_factor > 1，增加对比度，contrast_factor < 1，减少对比度
img_contrast = tf.clip_by_value(img_contrast1, 0, 1)

# 色调
img_hue = tf.image.adjust_hue(img, delta=0.1) #delta > 0，增加色调，delta < 0，减少色调
# 饱和度
img_saturation = tf.image.adjust_saturation(img, saturation_factor=2) #saturation_factor > 1，增加饱和度，saturation_factor < 1，减少饱和度
# rc =resize_center_crop(img)
# rp = resize_pad_crop(img)

rp = tf.image.resize_with_pad(img, 300, 300)
# g_noise = gaussian_noise(img, 0, 0.1)
# noise = salt_and_pepper_noise(img)
# mask_img = random_mask(img, mask_size=200, num_blocks=3)

display_images([img,rp])
display_images([img_flipped])
display_images([img_updown])
display_images([img_rotated])
display_images([img_resized])
display_images([img_cropped])
display_images([img_bright])
display_images([img_contrast])
display_images([img_hue])
display_images([img_saturation])


# display_images([[img], [img, img],[img, img, img]], "Original Image")
# display_images([[img],[img, img]], "Original Image")

#水平翻转
img_flipped = tf.image.flip_left_right(img)
#上下翻转
img_updown = tf.image.flip_up_down(img)
#旋转
img_rotated = tf.image.rot90(img, k=4)
#缩放
img_resized = tf.image.resize(img, size=[300, 300])
#按中心裁剪
img_cropped = tf.image.central_crop(img, central_fraction=0.5)
#按尺寸进行裁剪
img_cropped = tf.image.crop_to_bounding_box(img, offset_height=100, offset_width=100, target_height=200, target_width=200)
#随机裁剪，使用随机种子
img_cropped = tf.image.random_crop(img, size=[400, 400, 3], seed=1)

# 颜色相关的增强
# 增加亮度
img_bright = tf.image.adjust_brightness(img, delta=0.3) #delta > 0，它会增加亮度，delta < 0，它会减少亮度
# 增加对比度
img_contrast1 = tf.image.adjust_contrast(img, contrast_factor=2)# contrast_factor > 1，增加对比度，contrast_factor < 1，减少对比度
img_contrast = tf.clip_by_value(img_contrast1, 0, 1)

# 色调
img_hue = tf.image.adjust_hue(img, delta=0.1) #delta > 0，增加色调，delta < 0，减少色调
# 饱和度
img_saturation = tf.image.adjust_saturation(img, saturation_factor=2) #saturation_factor > 1，增加饱和度，saturation_factor < 1，减少饱和度
# rc =resize_center_crop(img)
# rp = resize_pad_crop(img)

def resize_center_crop(img, target_size=224):
    """
    保持宽高比调整大小后中心裁剪

    参数:
        img: 输入图像张量
        target_size: 目标尺寸(正方形)

    返回:
        中心裁剪后的图像
    """
    height, width, ch = img.shape
    aspect_ratio = width / height

    # 计算保持宽高比的临时尺寸
    if aspect_ratio > 1:  # 宽>高
        target_height = target_size
        target_width = int(round(target_size * aspect_ratio))
    else:  # 高>=宽
        target_width = target_size
        target_height = int(round(target_size / aspect_ratio))

    # 调整图像大小
    scale_img = tf.image.resize(img, [target_height, target_width])

    # 计算裁剪偏移量
    offset_height = (target_height - target_size) // 2
    offset_width = (target_width - target_size) // 2

    # 中心裁剪
    crop_img = tf.image.crop_to_bounding_box(
        scale_img,
        offset_height,
        offset_width,
        target_size,
        target_size
    )

    return crop_img


def resize_pad_crop(img, target_size=224):
    """
    保持宽高比调整大小后填充为正方形

    参数:
        img: 输入图像张量
        target_size: 目标尺寸(正方形)

    返回:
        填充后的图像
    """
    height, width, ch = img.shape
    aspect_ratio = width / height

    # 计算保持宽高比的临时尺寸
    if aspect_ratio > 1:  # 宽>高
        target_width = target_size
        target_height = int(round(target_size / aspect_ratio))
    else:  # 高>=宽
        target_height = target_size
        target_width = int(round(target_size * aspect_ratio))

    # 调整图像大小
    scale_img = tf.image.resize(img, [target_height, target_width])

    # 计算需要填充的像素数
    pad_height = target_size - target_height
    pad_width = target_size - target_width

    # 计算上下左右的填充量
    pad_top = pad_height // 2
    pad_bottom = pad_height - pad_top
    pad_left = pad_width // 2
    pad_right = pad_width - pad_left

    print(f"填充量: 上{pad_top}, 下{pad_bottom}, 左{pad_left}, 右{pad_right}")
    # 定义填充模式(只在高度和宽度维度填充)
    padding = [[pad_top, pad_bottom], [pad_left, pad_right], [0, 0]]
    # 执行填充(用0填充)
    pad_img = tf.pad(scale_img, padding, "CONSTANT", constant_values=0)

    return pad_img

def gaussian_noise(img, mean=0, stddev=0.1):
    """
    添加高斯噪声

    参数:
        img: 输入图像
        mean: 噪声均值
        stddev: 噪声标准差

    返回:
        添加噪声后的图像
    """
    # 生成高斯噪声
    noise = tf.random.normal(
        shape=tf.shape(img),
        mean=mean,
        stddev=stddev,
        dtype=tf.float32
    )
    # 添加噪声并裁剪到[0,1]范围
    noise_img = tf.clip_by_value(img + noise, 0, 1)
    return noise_img


def salt_and_pepper_noise(img, prob=0.1):
    """
    添加椒盐噪声

    参数:
        img: 输入图像
        prob: 噪声像素比例

    返回:
        添加噪声后的图像
    """
    height, width, chs = img.shape
    # 计算噪声像素数量
    num_noise_pixels = int(height * width * prob)

    # 随机生成噪声像素坐标
    y_coords = np.random.randint(0, height, num_noise_pixels)
    x_coords = np.random.randint(0, width, num_noise_pixels)

    # 随机决定是盐噪声(1)还是胡椒噪声(0)
    salt_or_pepper = np.random.choice([0, 1], size=num_noise_pixels * chs)
    salt_or_pepper = salt_or_pepper.reshape(num_noise_pixels, chs)

    # 创建图像副本并添加噪声
    noise_img = img.numpy().copy()
    for i in range(num_noise_pixels):
        noise_img[y_coords[i], x_coords[i]] = salt_or_pepper[i]

    return noise_img


def random_mask(img, mask_size=50, num_blocks=1):
    """
    随机添加黑色遮挡块

    参数:
        img: 输入图像
        mask_size: 遮挡块大小
        num_blocks: 遮挡块数量

    返回:
        添加遮挡后的图像
    """
    height, width, chs = img.shape
    img_np = img.numpy()
    # 添加指定数量的遮挡块
    for _ in range(num_blocks):
        # 随机生成遮挡块左上角坐标
        y = np.random.randint(0, height - mask_size)
        x = np.random.randint(0, width - mask_size)
        # 将区域设置为黑色(0)
        img_np[y:y + mask_size, x:x + mask_size] = 0

    return img_np

import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
from torchvision.io import read_image

import numpy as np
import matplotlib.pyplot as plt


def load_image(image_path):
    image = read_image(image_path) / 255.0
    return image


# image_list, [img1, img2, img3, ...]
# [[img1, img2], [img3, img4]]
# [[img1], [img2, img3]]
# [img1, [img2, img2]]
# [img1]
def display_images(image_list=None, title=None):
    if image_list is None:
        image_list = []

    if title is None:
        title = ""

    # (False, True)
    if not any(isinstance(i, list) for i in image_list):
        image_list = [image_list]  # [img, img, img] => [[img, img, img]]

    rows = len(image_list)
    cols = max(len(row) if isinstance(row, list) else 1 for row in image_list)

    plt.suptitle(title)
    fig, ax = plt.subplots(rows, cols)

    # 确保ax是2D数组
    # ax => [[ax]], [ax]=>[[ax]]
    ax = np.atleast_2d(ax)

    for i, row in enumerate(image_list):
        if not isinstance(row, list):
            row = [row]

        for j, img in enumerate(row):
            ax[i, j].imshow(img)
            ax[i, j].axis("off")

        for j in range(len(row), cols):
            ax[i, j].axis("off")

    # plt.tight_layout()
    plt.show()

image = load_image("./kpbl.PNG").permute(1, 2, 0)
print(image)
display_images(image, "Original Image")

# 基本变换

# tans_h_flip = transforms.RandomHorizontalFlip(p=1)
# img_h_flip = tans_h_flip(img).permute(1, 2, 0)
# tans_v_flip = transforms.RandomVerticalFlip(p=1)
# img_h_flip = tans_v_flip(img).permute(1, 2, 0)

# 缩放
# resize_img = transforms.Resize(size=(300, 300))
# img = resize_img(img).permute(1, 2, 0)

# 旋转
# rotate_img = transforms.RandomRotation(degrees=90, expand=True)
# img = rotate_img(img).permute(1, 2, 0)
# print(img_h_flip.shape)
# old_img = img.permute(1, 2, 0)

# 裁剪
# crop_img = transforms.CenterCrop(size=(300, 300))
# c_img = crop_img(img).permute(1, 2, 0)

# 随机裁剪
# random_crop = transforms.RandomCrop(size=(300, 300))
# c_img = random_crop(img).permute(1, 2, 0)

# 随机尺寸裁剪
#scale:其含义是：裁剪框的面积相对于原图像面积的比例范围
#ratio:其含义是：裁剪框的宽高比范围
# random_resized_crop = transforms.RandomResizedCrop(size=(300, 300), scale=(0.5, 1.0), ratio=(0.75, 1.333))
# c_img = random_resized_crop(img).permute(1, 2, 0)

def resize_center_crop(img, target_size=224):
    chs, height, width = img.shape 
    aspect_ratio = width / height

    if aspect_ratio > 1:
        target_height = target_size
        target_width = int(round(target_size * aspect_ratio))
    else:
        target_width = target_size
        target_height = int(round(target_size / aspect_ratio))
    
    resize_transform = transforms.Resize((target_height, target_width))
    resized_image = resize_transform(img)

    center_crop_transform = transforms.CenterCrop(target_size)
    cropped_image = center_crop_transform(resized_image)

    return cropped_image
    
def resize_pad_crop(img, target_size=224):
    chs, height, width = img.shape 
    aspect_ratio = width / height

    if aspect_ratio > 1:
        target_width = target_size
        target_height = int(round(target_size / aspect_ratio))
    else:
        target_height = target_size
        target_width = int(round(target_size * aspect_ratio))
    
    resize_transform = transforms.Resize((target_height, target_width)) 
    resized_image = resize_transform(img) 
    
    # Center crop 
    center_crop_transform = transforms.CenterCrop(target_size) 
    cropped_image = center_crop_transform(resized_image) 

    # RGB, RGBA,如果A全为0，表示透明，如果A全为255，表示不透明【0-255】
    #【0-1.0】，0表示示透明，1表示不透明
    if chs == 4:
        cropped_image[3,:,:]=1.0
    
    return cropped_image

# fivecrop
# five_crop = transforms.FiveCrop(size=(300, 300))
# ten crop
# ten_crop = transforms.TenCrop(size=(300, 300))
# c_img = ten_crop(img)
# img_list = []
# for i, img in enumerate(c_img):
#     img_list.append(img.permute(1, 2, 0))

# crop_img = resize_pad_crop(img).permute(1,2,0)
# gau_img = gaussian_noise(img).permute(1,2,0)
# sp_img = salt_and_pepper_noise(img)

import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"

from torchvision.datasets import VOCDetection
import torchvision.transforms as transforms


import numpy as np
import matplotlib.pyplot as plt

transform = transforms.Compose([
    transforms.ToTensor(),
])
#下载VOC 2007 数据集
train_data = VOCDetection(root='./datae',
                          year='2007',
                          image_set='train',
                          download=True,
                          transform=transform)


def display_images(image_list=None, title=None):
    if image_list is None:
        image_list = []

    if title is None:
        title = ""

    # (False, True)
    if not any(isinstance(i, list) for i in image_list):
        image_list = [image_list]  # [img, img, img] => [[img, img, img]]

    rows = len(image_list)
    cols = max(len(row) if isinstance(row, list) else 1 for row in image_list)

    plt.suptitle(title)
    fig, ax = plt.subplots(rows, cols)

    # 确保ax是2D数组
    # ax => [[ax]], [ax]=>[[ax]]
    ax = np.atleast_2d(ax)

    for i, row in enumerate(image_list):
        if not isinstance(row, list):
            row = [row]

        for j, img in enumerate(row):
            ax[i, j].imshow(img)
            ax[i, j].axis("off")

        for j in range(len(row), cols):
            ax[i, j].axis("off")

    # plt.tight_layout()
    plt.show()
    
images = []
for i in range(5):
    image, ann = train_data[i]
    images.append(image.permute(1,2,0))
    # print(image.shape, ann)
display_images(images)

import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
import tensorflow as tf

from torchvision.datasets import VOCDetection
import torchvision.transforms as transforms

import albumentations as A

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches

# 数据集根目录以及子目录定义
VOC_ROOT_DIR = './data/VOCdevkit/VOC2007/'
ANNOTATIONS_DIR = 'Annotations'
IMAGE_DIR = 'JPEGImages'

# 构建注释和图像目录路径
anno_dir = os.path.join(VOC_ROOT_DIR, ANNOTATIONS_DIR)  # 修正了变量名和拼写错误
image_dir = os.path.join(VOC_ROOT_DIR, IMAGE_DIR)  # 修正了变量名和拼写错误

# 打印路径以检查正确性
print(anno_dir)
print(image_dir)


def load_image(path):
    image = tf.io.read_file(path)
    decode_image = tf.image.decode_jpeg(image, channels=3)
    img = tf.image.convert_image_dtype(decode_image, tf.float32)

    return img


# %%
# image_list, [img1, img2, img3, ...]
# [[img1, img2], [img3, img4]]
# [[img1], [img2, img3]]
# [img1, [img2, img2]]
# [img1]
def display_images(image_list=None, title=None):
    if image_list is None:
        image_list = []

    if title is None:
        title = ""

    # (False, True)
    if not any(isinstance(i, list) for i in image_list):
        image_list = [image_list]  # [img, img, img] => [[img, img, img]]

    rows = len(image_list)
    cols = max(len(row) if isinstance(row, list) else 1 for row in image_list)

    plt.suptitle(title)
    fig, ax = plt.subplots(rows, cols)

    # 确保ax是2D数组
    # ax => [[ax]], [ax]=>[[ax]]
    ax = np.atleast_2d(ax)

    for i, row in enumerate(image_list):
        if not isinstance(row, list):
            row = [row]

        for j, img in enumerate(row):
            ax[i, j].imshow(img)
            ax[i, j].axis("off")

        for j in range(len(row), cols):
            ax[i, j].axis("off")

    plt.tight_layout()
    plt.show()


# %%
def resize_center_crop(img, target_size=224):
    height, width, ch = img.shape
    aspect_ratio = width / height

    if aspect_ratio > 1:
        target_height = target_size
        target_width = int(round(target_size * aspect_ratio))
    else:
        target_width = target_size
        target_height = int(round(target_size / aspect_ratio))

    scale_img = tf.image.resize(img, [target_height, target_width])

    offset_height = (target_height - target_size) // 2
    offset_width = (target_width - target_size) // 2

    crop_img = tf.image.crop_to_bounding_box(scale_img,
                                             offset_height,
                                             offset_width,
                                             target_size,
                                             target_size)

    return crop_img

img_path = os.path. join(image_dir,'000005.jpg')
print(img_path)

load_image(image_dir)
display_images([image_dir])

import tensorflow as tf
from tensorflow.keras import layers, models
from tensorflow.keras.optimizers import Adam

import pandas as pd
import numpy as np

from pathlib import Path

from utils import display_images

import albumentations as A

from sklearn.model_selection import train_test_split

BASE_PATH=Path("./banana-detection/banana-detection/bananas_train/")
IMAGES_PATH= BASE_PATH/"images"
LABEL_PATH = BASE_PATH/"label.csv"

print(IMAGES_PATH, LABEL_PATH)

#df(DataFrame)
df = pd.read_csv(LABEL_PATH)
filenames = df['img_name'].tolist()
print(filenames)

#通过该函数读取的图片都是归一化处理后的图片
def load_image(path):
    image = tf.io.read_file(path)
    decode_image = tf.image.decode_jpeg(image, channels=3)
    img = tf.image.convert_image_dtype(decode_image, tf.float32)

    return img

images = []
bbox_list = []
for i in range(5):
    if i == 4:
        image_path = IMAGES_PATH / filenames[i]
        img = load_image(str(image_path))
        images.append(img)

        bbox = [(df['xmin'][i],df['ymin'][i],df['xmax'][i],df['ymax'][i])]
        bbox_list.append(bbox)

#bboxes_list = [[(), ()...]]
display_images(images, bboxes_list=bbox_list)

def load_model(model_name, num_classes):
    if model_name == 'VGG16':
        # weights='imagenet'表示加载VGG16在ImageNet上预训练的权重
        # include_top=False表示不包含全连接层
        base_model = tf.keras.applications.VGG16(
            weights='imagenet',
            include_top=False,
            input_shape=(224, 224, 3)
        )
    elif model_name == 'ResNet50':
        base_model = tf.keras.applications.ResNet50(
            weights='imagenet',
            include_top=False,
            input_shape=(224, 224, 3)
        )

    # 冻结卷积层的参数
    base_model.trainable = False

    # 替换全连接层
    model = models.Sequential([
        base_model,
        layers.Flatten(),
        layers.Dense(128, activation='relu'),
        layers.Dense(64, activation='relu'),
        layers.Dense(32, activation='relu'),
        layers.Dense(num_classes, activation='sigmoid')
    ])
    model.summary()
    return model

model_name = 'VGG16'
num_classes = 4
#
model = load_model(model_name, num_classes)

# 导入必要的库
import tensorflow as tf
from tensorflow.keras import layers, models
from tensorflow.keras.optimizers import Adam
import pandas as pd
import numpy as np
from pathlib import Path
import albumentations as A
from sklearn.model_selection import train_test_split

# 定义基础路径和子目录
BASE_PATH = Path("./banana-detection/banana-detection/bananas_train/")
IMAGES_PATH = BASE_PATH / "images"
LABEL_PATH = BASE_PATH / "label.csv"

# 打印图像和标签文件夹路径以确认路径正确
print(IMAGES_PATH, LABEL_PATH)

# 从CSV文件中读取标签信息
df = pd.read_csv(LABEL_PATH)
filenames = df['img_name'].tolist()  # 获取所有图像文件名列表
print(filenames)

# 定义函数用于加载并归一化图像
def load_image(path):
    image = tf.io.read_file(path)  # 读取图像文件
    decode_image = tf.image.decode_jpeg(image, channels=3)  # 解码JPEG格式的图像
    img = tf.image.convert_image_dtype(decode_image, tf.float32)  # 将图像转换为浮点类型并归一化到[0,1]
    return img

# 初始化图像和边界框列表
images = []
bbox_list = []

# 加载前5个图像及其对应的边界框
for i in range(5):
    if i == 4:
        image_path = IMAGES_PATH / filenames[i]
        img = load_image(str(image_path))
        images.append(img)
        bbox = [(df['xmin'][i], df['ymin'][i], df['xmax'][i], df['ymax'][i])]
        bbox_list.append(bbox)

# 使用display_images函数显示图像和边界框（假设已定义）
# display_images(images, bboxes_list=bbox_list)

# 定义函数来加载预训练模型（VGG16或ResNet50）
def load_model(model_name, num_classes):
    if model_name == 'VGG16':
        base_model = tf.keras.applications.VGG16(
            weights='imagenet',
            include_top=False,
            input_shape=(224, 224, 3)
        )
    elif model_name == 'ResNet50':
        base_model = tf.keras.applications.ResNet50(
            weights='imagenet',
            include_top=False,
            input_shape=(224, 224, 3)
        )
    base_model.trainable = False  # 冻结卷积层参数
    
    # 添加自定义顶层
    model = models.Sequential([
        base_model,
        layers.Flatten(),
        layers.Dense(128, activation='relu'),
        layers.Dense(64, activation='relu'),
        layers.Dense(32, activation='relu'),
        layers.Dense(num_classes, activation='sigmoid')
    ])
    model.summary()
    return model

model_name = 'VGG16'
num_classes = 4
model = load_model(model_name, num_classes)

# 处理每行数据以提取图像和边界框
def process_row(row):
    image_path = IMAGES_PATH / row['img_name']
    image = load_image(str(image_path))
    bboxes = [(row['xmin'], row['ymin'], row['xmax'], row['ymax'])]
    return image, bboxes

image_bboxes = df.apply(lambda row: process_row(row), axis=1)
images, all_bboxes = zip(*image_bboxes)

# 数据增强函数
def data_augmentation(image, bboxes):
    transform = A.Compose([A.Resize(224, 224)], bbox_params=A.BboxParams(format='pascal_voc', label_fields=[]))
    transformed = transform(image=image.numpy(), bboxes=bboxes)
    return transformed['image'], transformed['bboxes']

# 归一化边界框
def normalize_bbox(bbox, image_width, image_height):
    x_min, y_min, x_max, y_max = bbox
    return [x_min / image_width, y_min / image_height, x_max / image_width, y_max / image_height]

# 应用数据增强和归一化
new_images, new_bbox_list = [], []
for i, img in enumerate(images):
    bboxes = all_bboxes[i]
    new_img, new_bboxes = data_augmentation(img, bboxes)
    new_images.append(new_img)
    bbox = new_bboxes[0]
    new_bbox = normalize_bbox(bbox, new_img.shape[1], new_img.shape[0])
    new_bbox_list.append(new_bbox)

# 转换为numpy数组并划分训练集和测试集
new_images = np.array(new_images, dtype=np.float32)
new_bbox_list = np.array(new_bbox_list, dtype=np.float32)
split = train_test_split(new_images, new_bbox_list, test_size=0.1, random_state=42)
train_images, test_images, train_labels, test_labels = split

# 编译模型
learning_rate = 1e-3
batch_size = 32
epochs = 25
opt = Adam(learning_rate=learning_rate)
model.compile(optimizer=opt, loss='mse', metrics=['accuracy'])

# 训练模型
H = model.fit(train_images, train_labels, validation_data=(test_images, test_labels),
              batch_size=batch_size, epochs=epochs, verbose=1)

# 绘制训练过程中的损失变化
plt.plot(range(epochs), H.history['loss'], label='train loss')
plt.plot(range(epochs), H.history['val_loss'], label='val loss')
plt.title('Loss')
plt.xlabel('Epoch #')
plt.ylabel('Loss')
plt.legend(loc='upper right')

# 保存模型
BASE_OUTPUT = Path("./output")
MODEL_PATH = BASE_OUTPUT / "single_obj_detect.keras"
model.save(MODEL_PATH)

# 加载模型并进行预测
model = tf.keras.models.load_model(MODEL_PATH, compile=False)
TEST_BASE_PATH = Path("./data/banana-detection/bananas_val")
TEST_IMAGE_PATH = TEST_BASE_PATH / "images"
test_image_path = TEST_IMAGE_PATH / "15.png"
test_image = load_image(str(test_image_path))
resized_image = tf.image.resize(test_image, (224, 224))
predict_image = tf.expand_dims(resized_image, axis=0)
preds = model.predict(predict_image)
(startX, startY, endX, endY) = preds[0]
startX, startY, endX, endY = int(startX * resized_image.shape[1]), int(startY * resized_image.shape[0]), int(endX * resized_image.shape[1]), int(endY * resized_image.shape[0])

# 显示预测结果
display_images(resized_image, bboxes_list=[[(startX, startY, endX, endY)]])

# 转换为numpy数组并划分训练集和测试集
new_images = np.array(new_images, dtype=np.float32)
new_bbox_list = np.array(new_bbox_list, dtype=np.float32)
split = train_test_split(new_images, new_bbox_list, test_size=0.1, random_state=42)
train_images, test_images, train_labels, test_labels = split

# 编译模型
learning_rate = 1e-3
batch_size = 32
epochs = 25
opt = Adam(learning_rate=learning_rate)
model.compile(optimizer=opt, loss='mse', metrics=['accuracy'])

# 训练模型
H = model.fit(train_images, train_labels, validation_data=(test_images, test_labels),
              batch_size=batch_size, epochs=epochs, verbose=1)

# 绘制训练过程中的损失变化
plt.plot(range(epochs), H.history['loss'], label='train loss')
plt.plot(range(epochs), H.history['val_loss'], label='val loss')
plt.title('Loss')
plt.xlabel('Epoch #')
plt.ylabel('Loss')
plt.legend(loc='upper right')

# 保存模型
BASE_OUTPUT = Path("./output")
MODEL_PATH = BASE_OUTPUT / "single_obj_detect.keras"
model.save(MODEL_PATH)

# 加载模型并进行预测
model = tf.keras.models.load_model(MODEL_PATH, compile=False)
TEST_BASE_PATH = Path("./data/banana-detection/bananas_val")
TEST_IMAGE_PATH = TEST_BASE_PATH / "images"
test_image_path = TEST_IMAGE_PATH / "15.png"
test_image = load_image(str(test_image_path))
resized_image = tf.image.resize(test_image, (224, 224))
predict_image = tf.expand_dims(resized_image, axis=0)
preds = model.predict(predict_image)
(startX, startY, endX, endY) = preds[0]
startX, startY, endX, endY = int(startX * resized_image.shape[1]), int(startY * resized_image.shape[0]), int(endX * resized_image.shape[1]), int(endY * resized_image.shape[0])

# 显示预测结果
display_images(resized_image, bboxes_list=[[(startX, startY, endX, endY)]])

TEST_BASE_PATH = Path("./data/banana-detection/bananas_val")
TEST_IMAGE_PATH = TEST_BASE_PATH / "images"
test_image_path = TEST_IMAGE_PATH / "15.png"
print(test_image_path)
test_image = load_image(str(test_image_path))

resized_image = tf.image.resize(test_image, (224, 224))

predict_image = tf.expand_dims(resized_image, axis=0)

#predict要求的输出是【batch_size, (height, width, channels)】
preds = model.predict(predict_image)
(startX, startY, endX, endY) = preds[0]

startX = int(startX * resized_image.shape[1])
startY = int(startY * resized_image.shape[0])
endX = int(endX * resized_image.shape[1])
endY = int(endY * resized_image.shape[0])

display_images(resized_image, bboxes_list=[[(startX, startY, endX, endY)]])

# 设置环境变量以避免潜在的库冲突
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"

# 导入必要的PyTorch模块
import torch
import torch.nn as nn  # 神经网络模块
import torch.optim as optim  # 优化算法

# 导入预训练模型及其权重
from torchvision import models
from torchvision.models import VGG16_Weights, ResNet18_Weights

# 数据加载工具
from torch.utils.data import DataLoader, Dataset

# 路径处理
from pathlib import Path

# 数据处理
import pandas as pd

# 图像读取
from torchvision.io import read_image

# 图像变换和增强
import albumentations as A
from albumentations.pytorch import ToTensorV2

def normalize_bbox(bbox, image_width, image_height):
    """
    将边界框坐标归一化到[0,1]范围（相对于图像尺寸）
    
    参数:
        bbox: 边界框坐标 (x_min, y_min, x_max, y_max)
        image_width: 图像宽度
        image_height: 图像高度
        
    返回:
        归一化后的边界框坐标 [x_min, y_min, x_max, y_max]
    """
    x_min, y_min, x_max, y_max = bbox
    x_min = x_min / image_width
    x_max = x_max / image_width
    y_min = y_min / image_height
    y_max = y_max / image_height
    return [x_min, y_min, x_max, y_max]

# 使用Albumentations定义图像变换
transform = A.Compose([
    A.Resize(224, 224),  # 将图像调整为224x224（许多CNN模型的标准输入尺寸）
    ToTensorV2()  # 转换为PyTorch张量
], bbox_params=A.BboxParams(
    format='pascal_voc',  # 边界框格式 (xmin, ymin, xmax, ymax)
    label_fields=[]  # 不需要额外的标签字段
))

class CustomDataset(Dataset):
    """
    自定义数据集类，用于加载图像和边界框
    
    参数:
        root_dir: 包含图像的根目录
        csv_file: 包含标注的CSV文件（img_name, label, xmin, ymin, xmax, ymax）
        transform: 要应用的图像变换
    """
    def __init__(self, root_dir, csv_file, transform=None):
        self.root_dir = root_dir  # 图像根目录
        self.annotations = pd.read_csv(csv_file)  # 读取标注CSV文件
        self.transform = transform  # 图像变换

    def __len__(self):
        """返回数据集中的样本数量"""
        return len(self.annotations)

    def __getitem__(self, idx):
        """
        获取单个样本
        
        参数:
            idx: 样本索引
            
        返回:
            image: 图像张量
            bbox: 归一化后的边界框坐标
        """
        # 获取图像文件名并构建完整路径
        img_name = self.annotations.iloc[idx, 0]
        img_path = Path(self.root_dir) / img_name
        
        # 读取图像并归一化到[0,1]范围
        image = read_image(str(img_path)) / 255.0

        # 获取边界框坐标并归一化
        bbox = self.annotations.iloc[idx, 2:].values.astype(float)
        bbox = normalize_bbox(bbox, image.shape[2], image.shape[1])
        bboxes = [bbox]  # 转换为列表形式（Albumentations要求）

        # 应用图像变换（如果有）
        if self.transform:
            transformed = self.transform(
                image=image.permute(1, 2, 0).numpy(),  # 将CHW转为HWC并转为numpy
                bboxes=bboxes
            )
            image = transformed['image']  # 变换后的图像
            bboxes = transformed['bboxes']  # 变换后的边界框

        return image, torch.tensor(bboxes[0], dtype=torch.float32)

class CustomModel(nn.Module):
    """
    自定义模型类，基于预训练模型构建
    
    参数:
        model_name: 模型名称（'vgg16'或'resnet18'）
        num_classes: 输出类别数（这里用于边界框回归，输出4个坐标值）
    """
    def __init__(self, model_name, num_classes):
        super(CustomModel, self).__init__()

        if model_name == 'vgg16':
            # 加载预训练的VGG16模型
            self.model = models.vgg16(weights=VGG16_Weights.IMAGENET1K_V1)

            # 冻结特征提取层的参数（不更新权重）
            for param in self.model.features.parameters():
                param.requires_grad = False

            # 计算特征维度
            with torch.no_grad():
                features = self.model.features(torch.randn(1, 3, 224, 224))
                features_dim = features.view(features.size(0), -1).size(1)

            # 替换分类器部分
            self.model.classifier = nn.Sequential(
                nn.Linear(features_dim, 128),
                nn.ReLU(),
                nn.Linear(128, 64),
                nn.ReLU(),
                nn.Linear(64, 32),
                nn.ReLU(),
                nn.Linear(32, num_classes),  # 输出4个坐标值
                nn.Sigmoid()  # 将输出限制在[0,1]范围
            )
            
        elif model_name == 'resnet18':
            # 加载预训练的ResNet18模型
            self.model = models.resnet18(weights=ResNet18_Weights.DEFAULT)

            # 冻结除全连接层外的所有参数
            for name, param in self.model.named_parameters():
                if not isinstance(param, nn.Linear):
                    param.requires_grad = False

            # 获取全连接层的输入特征数
            in_features = self.model.fc.in_features

            # 替换全连接层
            self.model.fc = nn.Sequential(
                nn.Linear(in_features, 128),
                nn.ReLU(),
                nn.Linear(128, 64),
                nn.ReLU(),
                nn.Linear(64, 32),
                nn.ReLU(),
                nn.Linear(32, num_classes),  # 输出4个坐标值
                nn.Sigmoid()  # 将输出限制在[0,1]范围
            )
        else:
            raise ValueError('未知模型名称: {}'.format(model_name))

        print(self.model)  # 打印模型结构

    def forward(self, x):
        """前向传播"""
        return self.model(x)

# 创建训练数据集和数据加载器
dataset = CustomDataset(
    root_dir='./data/banana-detection/bananas_train/images/',
    csv_file='./data/banana-detection/bananas_train/label.csv',
    transform=transform
)

# 批量加载数据（批量大小32，打乱顺序）
dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# 检测并设置设备（优先使用GPU，其次是MPS，最后是CPU）
if torch.cuda.is_available():
    device = torch.device("cuda")
elif torch.backends.mps.is_available():
    device = torch.device("mps")
else:
    device = torch.device("cpu")

print(f"使用设备: {device}")

# 模型配置
model_name = 'vgg16'  # 选择模型架构
num_classes = 4  # 输出4个坐标值（xmin, ymin, xmax, ymax）
model = CustomModel(model_name=model_name, num_classes=num_classes)
model.to(device)  # 将模型移动到指定设备

# 定义损失函数（均方误差损失）和优化器（Adam）
loss_fn = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# 训练参数
num_epochs = 25  # 训练轮数

# 训练循环
for epoch in range(num_epochs):
    for images, bboxes in dataloader:
        # 将数据移动到指定设备
        images, bboxes = images.to(device), bboxes.to(device)

        # 前向传播
        outputs = model(images)  # 预测边界框
        loss = loss_fn(outputs, bboxes)  # 计算损失

        # 反向传播和优化
        optimizer.zero_grad()  # 清空梯度
        loss.backward()  # 反向传播
        optimizer.step()  # 更新参数

    # 打印每个epoch的损失
    print(f'轮次 [{epoch+1}/{num_epochs}], 损失: {loss.item():.4f}')

# 创建输出目录
output_dir = Path('./output')
output_dir.mkdir(exist_ok=True)  # 如果目录不存在则创建

# 保存模型权重
model_save_path = output_dir/'single_obj_detect.pth'
torch.save(model.state_dict(), model_save_path)
print(f"模型已保存到: {model_save_path}")

# 测试模型效果
from utils import display_images  # 导入可视化工具

# 加载测试图像
test_image_path = './data/banana-detection/bananas_val/images/0.png'
test_image = read_image(test_image_path) / 255.0  # 读取并归一化
test_image = transform(image=test_image.permute(1, 2, 0).numpy())['image']  # 应用变换

# 准备模型输入
model.eval()  # 设置为评估模式
pre_image = test_image.unsqueeze(0).to(device)  # 添加批次维度并移动到设备

# 进行预测
with torch.no_grad():
    outputs = model(pre_image)
    (startX, startY, endX, endY) = outputs[0].cpu().numpy()

# 将归一化坐标转换回图像尺寸
height, width = test_image.shape[1], test_image.shape[2]
startX = int(startX * width)
startY = int(startY * height)
endX = int(endX * width)
endY = int(endY * height)

# 可视化结果
display_images(
    test_image.permute(1, 2, 0),  # 将CHW转为HWC
    bboxes_list=[[(startX, startY, endX, endY)]]  # 预测的边界框
)