# python test 1G - calculate warmth
import cv2 as cv
import os
from glob import glob
import csv
import numpy

#baseDir = "C:/Users/me/sd-web/webui/outputs/txt2img-images/"
#workingDir = "2024-01-03"
baseDir = "" #use cwd
workingDir = "images"

detector = cv.FaceDetectorYN.create("face_detection_yunet_2023mar.onnx",  "", (0, 0))

def imgAnalysis (fn):
	img = cv.imread(fn)
	imgW = img.shape[0]
	imgH = img.shape[1]

	detector.setInputSize((imgW, imgH))
	_, faces = detector.detect(img)

	faceArea = 0.0
	if faces is not None: 
		for face in faces:
			coords = face[:-1].astype(numpy.int32)
			#we may double account if we find overlapping boxes that bound multiple faces
			faceArea += coords[2] * coords[3]
	facialOccupancy = faceArea / (imgW * imgH)
	
	img = cv.cvtColor(img,cv.COLOR_BGR2HLS)
	lumIdx = 1
	satIdx = 2
	hueIdx = 0
	
	#convert hue to warm scale
	w = numpy.int16(img[:,:,hueIdx])
	w[:,:] -= 15	# rotate hues 30deg counter clockwise, warmest at 0
	w[:,:] %= 180	# shift scale to 0..179
	w[:,:] -= 90
	w[:,:] = abs(w[:,:])
	img[:,:,hueIdx] = numpy.uint8(w)
	
	means = img.mean ((0,1))
	stds = img.std ((0,1))
	
	result = [facialOccupancy * 100.0]
	result.append(means[lumIdx])
	result.append(stds[lumIdx])
	result.append(means[satIdx])
	result.append(means[hueIdx])
	return (result)

#create a list of unique seed numbers  
seeds = []
#and a list of file indexes for that seed
fIndexes = []
files = glob(baseDir+workingDir+"/*.png")
for f in files:
	fn = os.path.basename(os.path.splitext(f)[0])
	fileIdx,seed = fn.split('-')
	if seed not in seeds:
		seeds.append(seed)
		fIndexes.append([])
	idx = seeds.index(seed)
	fIndexes[idx].append(fileIdx)

results = []

#create the csv file header row
#honestly I oughta use a dictionary here I suppose
#but this is a pretty quick'n'dirty implementation all round!
head = ["Seed"]
head.append("Cntrl fIdx")
head.append("Face")
head.append("LumMean")
head.append("LumSD")
head.append("SatMean")
head.append("WarmthMean")
head.append("Var1 fIdx")
head.append("Face")
head.append("LumMean")
head.append("LumSD")
head.append("SatMean")
head.append("WarmthMean")
head.append("Var2 fIdx")
head.append("Face")
head.append("LumMean")
head.append("LumSD")
head.append("SatMean")
head.append("WarmthMean")

results.append(head)

for s in seeds:
	idx = seeds.index(s)
	result = [s]
	for f in fIndexes[idx]:
		result.append(f)
		fn = baseDir+workingDir+"/"+f+"-"+s+".png"
		analr = imgAnalysis(fn)
		result.extend(analr)
	print (result)
	results.append(result)

with open(workingDir+".csv",'w', newline='') as myfile:
	wr = csv.writer(myfile, quoting=csv.QUOTE_NONNUMERIC)
	wr.writerows(results)



