summaryrefslogtreecommitdiff
path: root/utils/raspberrypi/ctt/ctt_awb.py
diff options
context:
space:
mode:
Diffstat (limited to 'utils/raspberrypi/ctt/ctt_awb.py')
-rw-r--r--utils/raspberrypi/ctt/ctt_awb.py118
1 files changed, 59 insertions, 59 deletions
diff --git a/utils/raspberrypi/ctt/ctt_awb.py b/utils/raspberrypi/ctt/ctt_awb.py
index d3130612..9a92033b 100644
--- a/utils/raspberrypi/ctt/ctt_awb.py
+++ b/utils/raspberrypi/ctt/ctt_awb.py
@@ -12,7 +12,7 @@ from scipy.optimize import fmin
"""
obtain piecewise linear approximation for colour curve
"""
-def awb(Cam,cal_cr_list,cal_cb_list,plot):
+def awb(Cam, cal_cr_list, cal_cb_list, plot):
imgs = Cam.imgs
"""
condense alsc calibration tables into one dictionary
@@ -21,7 +21,7 @@ def awb(Cam,cal_cr_list,cal_cb_list,plot):
colour_cals = None
else:
colour_cals = {}
- for cr,cb in zip(cal_cr_list,cal_cb_list):
+ for cr, cb in zip(cal_cr_list, cal_cb_list):
cr_tab = cr['table']
cb_tab = cb['table']
"""
@@ -29,7 +29,7 @@ def awb(Cam,cal_cr_list,cal_cb_list,plot):
"""
cr_tab= cr_tab/np.min(cr_tab)
cb_tab= cb_tab/np.min(cb_tab)
- colour_cals[cr['ct']] = [cr_tab,cb_tab]
+ colour_cals[cr['ct']] = [cr_tab, cb_tab]
"""
obtain data from greyscale macbeth patches
"""
@@ -42,17 +42,17 @@ def awb(Cam,cal_cr_list,cal_cb_list,plot):
Note: if alsc is disabled then colour_cals will be set to None and the
function will just return the greyscale patches
"""
- r_patchs,b_patchs,g_patchs = get_alsc_patches(Img,colour_cals)
+ r_patchs, b_patchs, g_patchs = get_alsc_patches(Img, colour_cals)
"""
- calculate ratio of r,b to g
+ calculate ratio of r, b to g
"""
r_g = np.mean(r_patchs/g_patchs)
b_g = np.mean(b_patchs/g_patchs)
- Cam.log += '\n r : {:.4f} b : {:.4f}'.format(r_g,b_g)
+ Cam.log += '\n r : {:.4f} b : {:.4f}'.format(r_g, b_g)
"""
The curve tends to be better behaved in so-called hatspace.
- R,B,G represent the individual channels. The colour curve is plotted in
- r,b space, where:
+ R, B, G represent the individual channels. The colour curve is plotted in
+ r, b space, where:
r = R/G
b = B/G
This will be referred to as dehatspace... (sorry)
@@ -71,18 +71,18 @@ def awb(Cam,cal_cr_list,cal_cb_list,plot):
"""
r_g_hat = r_g/(1+r_g+b_g)
b_g_hat = b_g/(1+r_g+b_g)
- Cam.log += '\n r_hat : {:.4f} b_hat : {:.4f}'.format(r_g_hat,b_g_hat)
- rbs_hat.append((r_g_hat,b_g_hat,Img.col))
- rb_raw.append((r_g,b_g))
+ Cam.log += '\n r_hat : {:.4f} b_hat : {:.4f}'.format(r_g_hat, b_g_hat)
+ rbs_hat.append((r_g_hat, b_g_hat, Img.col))
+ rb_raw.append((r_g, b_g))
Cam.log += '\n'
Cam.log += '\nFinished processing images'
"""
sort all lits simultaneously by r_hat
"""
- rbs_zip = list(zip(rbs_hat,rb_raw))
- rbs_zip.sort(key=lambda x:x[0][0])
- rbs_hat,rb_raw = list(zip(*rbs_zip))
+ rbs_zip = list(zip(rbs_hat, rb_raw))
+ rbs_zip.sort(key=lambda x: x[0][0])
+ rbs_hat, rb_raw = list(zip(*rbs_zip))
"""
unzip tuples ready for processing
"""
@@ -91,7 +91,7 @@ def awb(Cam,cal_cr_list,cal_cb_list,plot):
"""
fit quadratic fit to r_g hat and b_g_hat
"""
- a,b,c = np.polyfit(rbs_hat[0],rbs_hat[1],2)
+ a, b, c = np.polyfit(rbs_hat[0], rbs_hat[1], 2)
Cam.log += '\nFit quadratic curve in hatspace'
"""
the algorithm now approximates the shortest distance from each point to the
@@ -113,11 +113,11 @@ def awb(Cam,cal_cr_list,cal_cb_list,plot):
def f(x):
return a*x**2 + b*x + c
"""
- iterate over points (R,B are x and y coordinates of points) and calculate
+ iterate over points (R, B are x and y coordinates of points) and calculate
distance to line in dehatspace
"""
dists = []
- for i, (R,B) in enumerate(zip(rbs_hat[0],rbs_hat[1])):
+ for i, (R, B) in enumerate(zip(rbs_hat[0], rbs_hat[1])):
"""
define function to minimise as square distance between datapoint and
point on curve. Squaring is monotonic so minimising radius squared is
@@ -129,7 +129,7 @@ def awb(Cam,cal_cr_list,cal_cb_list,plot):
"""
perform optimisation with scipy.optmisie.fmin
"""
- x_hat = fmin(f_min,R,disp=0)[0]
+ x_hat = fmin(f_min, R, disp=0)[0]
y_hat = f(x_hat)
"""
dehat
@@ -179,16 +179,16 @@ def awb(Cam,cal_cr_list,cal_cb_list,plot):
"""
r_fit = r_hat_fit/(1-r_hat_fit-b_hat_fit)
b_fit = b_hat_fit/(1-r_hat_fit-b_hat_fit)
- c_fit = np.round(rbs_hat[2],0)
+ c_fit = np.round(rbs_hat[2], 0)
"""
round to 4dp
"""
- r_fit = np.where((1000*r_fit)%1<=0.05,r_fit+0.0001,r_fit)
- r_fit = np.where((1000*r_fit)%1>=0.95,r_fit-0.0001,r_fit)
- b_fit = np.where((1000*b_fit)%1<=0.05,b_fit+0.0001,b_fit)
- b_fit = np.where((1000*b_fit)%1>=0.95,b_fit-0.0001,b_fit)
- r_fit = np.round(r_fit,4)
- b_fit = np.round(b_fit,4)
+ r_fit = np.where((1000*r_fit)%1<=0.05, r_fit+0.0001, r_fit)
+ r_fit = np.where((1000*r_fit)%1>=0.95, r_fit-0.0001, r_fit)
+ b_fit = np.where((1000*b_fit)%1<=0.05, b_fit+0.0001, b_fit)
+ b_fit = np.where((1000*b_fit)%1>=0.95, b_fit-0.0001, b_fit)
+ r_fit = np.round(r_fit, 4)
+ b_fit = np.round(b_fit, 4)
"""
The following code ensures that colour temperature decreases with
increasing r/g
@@ -200,8 +200,8 @@ def awb(Cam,cal_cr_list,cal_cb_list,plot):
while i > 0 :
if c_fit[i] > c_fit[i-1]:
Cam.log += '\nColour temperature increase found\n'
- Cam.log += '{} K at r = {} to '.format(c_fit[i-1],r_fit[i-1])
- Cam.log += '{} K at r = {}'.format(c_fit[i],r_fit[i])
+ Cam.log += '{} K at r = {} to '.format(c_fit[i-1], r_fit[i-1])
+ Cam.log += '{} K at r = {}'.format(c_fit[i], r_fit[i])
"""
if colour temperature increases then discard point furthest from
the transformed fit (dehatspace)
@@ -209,8 +209,8 @@ def awb(Cam,cal_cr_list,cal_cb_list,plot):
error_1 = abs(dists[i-1])
error_2 = abs(dists[i])
Cam.log += '\nDistances from fit:\n'
- Cam.log += '{} K : {:.5f} , '.format(c_fit[i],error_1)
- Cam.log += '{} K : {:.5f}'.format(c_fit[i-1],error_2)
+ Cam.log += '{} K : {:.5f} , '.format(c_fit[i], error_1)
+ Cam.log += '{} K : {:.5f}'.format(c_fit[i-1], error_2)
"""
find bad index
note that in python false = 0 and true = 1
@@ -221,9 +221,9 @@ def awb(Cam,cal_cr_list,cal_cb_list,plot):
"""
delete bad point
"""
- r_fit = np.delete(r_fit,bad)
- b_fit = np.delete(b_fit,bad)
- c_fit = np.delete(c_fit,bad).astype(np.uint16)
+ r_fit = np.delete(r_fit, bad)
+ b_fit = np.delete(b_fit, bad)
+ c_fit = np.delete(c_fit, bad).astype(np.uint16)
"""
note that if a point has been discarded then the length has decreased
by one, meaning that decreasing the index by one will reassess the kept
@@ -235,7 +235,7 @@ def awb(Cam,cal_cr_list,cal_cb_list,plot):
"""
return formatted ct curve, ordered by increasing colour temperature
"""
- ct_curve = list(np.array(list(zip(b_fit,r_fit,c_fit))).flatten())[::-1]
+ ct_curve = list(np.array(list(zip(b_fit, r_fit, c_fit))).flatten())[::-1]
Cam.log += '\nFinal CT curve:'
for i in range(len(ct_curve)//3):
j = 3*i
@@ -247,15 +247,15 @@ def awb(Cam,cal_cr_list,cal_cb_list,plot):
plotting code for debug
"""
if plot:
- x = np.linspace(np.min(rbs_hat[0]),np.max(rbs_hat[0]),100)
+ x = np.linspace(np.min(rbs_hat[0]), np.max(rbs_hat[0]), 100)
y = a*x**2 + b*x + c
- plt.subplot(2,1,1)
+ plt.subplot(2, 1, 1)
plt.title('hatspace')
- plt.plot(rbs_hat[0],rbs_hat[1],ls='--',color='blue')
- plt.plot(x,y,color='green',ls='-')
- plt.scatter(rbs_hat[0],rbs_hat[1],color='red')
+ plt.plot(rbs_hat[0], rbs_hat[1], ls='--', color='blue')
+ plt.plot(x, y, color='green', ls='-')
+ plt.scatter(rbs_hat[0], rbs_hat[1], color='red')
for i, ct in enumerate(rbs_hat[2]):
- plt.annotate(str(ct),(rbs_hat[0][i],rbs_hat[1][i]))
+ plt.annotate(str(ct), (rbs_hat[0][i], rbs_hat[1][i]))
plt.xlabel('$\hat{r}$')
plt.ylabel('$\hat{b}$')
"""
@@ -265,13 +265,13 @@ def awb(Cam,cal_cr_list,cal_cb_list,plot):
# ax = plt.gca()
# ax.set_aspect('equal')
plt.grid()
- plt.subplot(2,1,2)
+ plt.subplot(2, 1, 2)
plt.title('dehatspace - indoors?')
- plt.plot(r_fit,b_fit,color='blue')
- plt.scatter(rb_raw[0],rb_raw[1],color='green')
- plt.scatter(r_fit,b_fit,color='red')
- for i,ct in enumerate(c_fit):
- plt.annotate(str(ct),(r_fit[i],b_fit[i]))
+ plt.plot(r_fit, b_fit, color='blue')
+ plt.scatter(rb_raw[0], rb_raw[1], color='green')
+ plt.scatter(r_fit, b_fit, color='red')
+ for i, ct in enumerate(c_fit):
+ plt.annotate(str(ct), (r_fit[i], b_fit[i]))
plt.xlabel('$r$')
plt.ylabel('$b$')
"""
@@ -286,15 +286,15 @@ def awb(Cam,cal_cr_list,cal_cb_list,plot):
"""
end of plotting code
"""
- return(ct_curve,np.round(transverse_pos,5),np.round(transverse_neg,5))
+ return(ct_curve, np.round(transverse_pos, 5), np.round(transverse_neg, 5))
"""
obtain greyscale patches and perform alsc colour correction
"""
-def get_alsc_patches(Img,colour_cals,grey=True):
+def get_alsc_patches(Img, colour_cals, grey=True):
"""
get patch centre coordinates, image colour and the actual
- patches for each channel,remembering to subtract blacklevel
+ patches for each channel, remembering to subtract blacklevel
If grey then only greyscale patches considered
"""
if grey:
@@ -316,12 +316,12 @@ def get_alsc_patches(Img,colour_cals,grey=True):
g_patchs = (patches[1]+patches[2])/2 - Img.blacklevel_16
if colour_cals == None:
- return r_patchs,b_patchs,g_patchs
+ return r_patchs, b_patchs, g_patchs
"""
find where image colour fits in alsc colour calibration tables
"""
cts = list(colour_cals.keys())
- pos = bisect_left(cts,col)
+ pos = bisect_left(cts, col)
"""
if img colour is below minimum or above maximum alsc calibration colour, simply
pick extreme closest to img colour
@@ -343,32 +343,32 @@ def get_alsc_patches(Img,colour_cals,grey=True):
bef_tabs = np.array(colour_cals[bef])
aft_tabs = np.array(colour_cals[aft])
col_tabs = (bef_tabs*db + aft_tabs*da)/(da+db)
- col_tabs = np.reshape(col_tabs,(2,12,16))
+ col_tabs = np.reshape(col_tabs, (2, 12, 16))
"""
- calculate dx,dy used to calculate alsc table
+ calculate dx, dy used to calculate alsc table
"""
- w,h = Img.w/2,Img.h/2
- dx,dy = int(-(-(w-1)//16)),int(-(-(h-1)//12))
+ w, h = Img.w/2, Img.h/2
+ dx, dy = int(-(-(w-1)//16)), int(-(-(h-1)//12))
"""
make list of pairs of gains for each patch by selecting the correct value
in alsc colour calibration table
"""
patch_gains = []
for cen in cen_coords:
- x,y = cen[0]//dx, cen[1]//dy
+ x, y = cen[0]//dx, cen[1]//dy
# We could probably do with some better spatial interpolation here?
- col_gains = (col_tabs[0][y][x],col_tabs[1][y][x])
+ col_gains = (col_tabs[0][y][x], col_tabs[1][y][x])
patch_gains.append(col_gains)
"""
multiply the r and b channels in each patch by the respective gain, finally
performing the alsc colour correction
"""
- for i,gains in enumerate(patch_gains):
+ for i, gains in enumerate(patch_gains):
r_patchs[i] = r_patchs[i] * gains[0]
b_patchs[i] = b_patchs[i] * gains[1]
"""
- return greyscale patches, g channel and correct r,b channels
+ return greyscale patches, g channel and correct r, b channels
"""
- return r_patchs,b_patchs,g_patchs
+ return r_patchs, b_patchs, g_patchs