import cv2 as cv
import matplotlib.pyplot as plt
import numpy as np
import copy
import os
import math
import pandas as pd

def dim3to2(rgb):
    const_val = math.sqrt(2)/2
    return np.dot(rgb[...,:3],np.array([[0,const_val,-const_val], [const_val,-const_val,0]]).T)

path = 'D:\\temp\\B题全部数据\\Rock'
filenames = os.listdir(path)

pct_list = pd.DataFrame(columns=['file name', 'oil percentage'])
for filename in filenames:
    sp_list = filename.split('.')
    if not sp_list[-1] == 'bmp':
        continue
    first = sp_list[0]#1-2
    type = first.split('-')[1]#2
    if type != '2':#只有2是荧光图
        continue

    fig = plt.imread(path + '\\' + filename)
    b, g, r = cv.split(fig)
    img = cv.merge([r, g, b])  # 为了适应fc中元组顺序为RGB，故调整
    res = dim3to2(img)

    width = res.shape[0]
    height = res.shape[1]
    cac_pct = 0#统计绿色和黄色

    green_yellow = copy.deepcopy(img)
    for x in range(width):
        for y in range(height):
            if (res[x, y, 0]**2 + res[x, y, 1]**2) > 30**2 and res[x, y, 0] > res[x, y, 1]:
                cac_pct += 1
            # if res[x, y, 0] > -30 and res[x, y, 1] < -25:#green
            #     cac_pct += 1
            # elif res[x, y, 0] >=40 and res[x, y, 1] >= -25:#yellow
            #     cac_pct += 1

    percent = cac_pct / (width * height)
    pct_list = pct_list.append({'file name': filename, 'oil percentage': percent}, ignore_index=True)
    print('filename=', filename, ' percentage=', percent)

pct_list.to_csv('D:\\其他\\data\\rock_oil_bmp.csv')