import random
import torch
import numpy as np
import os
import pandas as pd
import cv2
from config import *

# 设置随机种子参数
def setup_seed(seed):
    np.random.seed(seed)
    random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

# 计算图像的像素平均值
def caculate_mean(img_dir, img_size):
    sum_r = 0
    sum_g = 0
    sum_b = 0
    count = 0

    for img_name in os.listdir(img_dir):
        img_path = os.path.join(img_dir, img_name)
        img = cv2.imread(img_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = cv2.resize(img, (img_size, img_size))
        sum_r += img[:, :, 0].mean()
        sum_g += img[:, :, 1].mean()
        sum_b += img[:, :, 2].mean()
        count += 1

    sum_r /= count
    sum_g /= count
    sum_b /= count
    img_mean = [sum_r/255., sum_g/255., sum_b/255.]
    return img_mean

# 接触角测量误差保存
def Save_data(data, root):
    data1 = np.array(data)
    df = pd.DataFrame({
        '迭代次数': pd.Series(np.arange(1, len(data) + 1, 1), index=np.arange(1, len(data) + 1, 1)),
        '损失数': pd.Series(data1, np.arange(1, len(data) + 1, 1))
    })
    df.to_csv(root)

if __name__ == '__main__':
    print(caculate_mean(img_dir=img_root, img_size=224))