538 lines
20 KiB
Python
538 lines
20 KiB
Python
# SPDX-License-Identifier: BSD-2-Clause
|
|
#
|
|
# Copyright (C) 2019, Raspberry Pi Ltd
|
|
# Copyright (C) 2024, Ideas on Board Oy
|
|
#
|
|
# Locate and extract Macbeth charts from images
|
|
# (Copied from: ctt_macbeth_locator.py)
|
|
|
|
# \todo Add debugging
|
|
|
|
import cv2
|
|
import os
|
|
from pathlib import Path
|
|
import numpy as np
|
|
import warnings
|
|
import logging
|
|
from sklearn import cluster as cluster
|
|
|
|
from .ctt_ransac import get_square_verts, get_square_centres
|
|
from .image import Image
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class MacbethError(Exception):
|
|
pass
|
|
|
|
|
|
# Reshape image to fixed width without distorting returns image and scale
|
|
# factor
|
|
def reshape(img, width):
|
|
factor = width / img.shape[0]
|
|
return cv2.resize(img, None, fx=factor, fy=factor), factor
|
|
|
|
|
|
# Correlation function to quantify match
|
|
def correlate(im1, im2):
|
|
f1 = im1.flatten()
|
|
f2 = im2.flatten()
|
|
cor = np.corrcoef(f1, f2)
|
|
return cor[0][1]
|
|
|
|
|
|
# @brief Compute coordinates of macbeth chart vertices and square centres
|
|
# @return (max_cor, best_map_col_norm, fit_coords, success)
|
|
#
|
|
# Also returns an error/success message for debugging purposes. Additionally,
|
|
# it scores the match with a confidence value.
|
|
#
|
|
# Brief explanation of the macbeth chart locating algorithm:
|
|
# - Find rectangles within image
|
|
# - Take rectangles within percentage offset of median perimeter. The
|
|
# assumption is that these will be the macbeth squares
|
|
# - For each potential square, find the 24 possible macbeth centre locations
|
|
# that would produce a square in that location
|
|
# - Find clusters of potential macbeth chart centres to find the potential
|
|
# macbeth centres with the most votes, i.e. the most likely ones
|
|
# - For each potential macbeth centre, use the centres of the squares that
|
|
# voted for it to find macbeth chart corners
|
|
# - For each set of corners, transform the possible match into normalised
|
|
# space and correlate with a reference chart to evaluate the match
|
|
# - Select the highest correlation as the macbeth chart match, returning the
|
|
# correlation as the confidence score
|
|
#
|
|
# \todo Clean this up
|
|
def get_macbeth_chart(img, ref_data):
|
|
ref, ref_w, ref_h, ref_corns = ref_data
|
|
|
|
# The code will raise and catch a MacbethError in case of a problem, trying
|
|
# to give some likely reasons why the problem occured, hence the try/except
|
|
try:
|
|
# Obtain image, convert to grayscale and normalise
|
|
src = img
|
|
src, factor = reshape(src, 200)
|
|
original = src.copy()
|
|
a = 125 / np.average(src)
|
|
src_norm = cv2.convertScaleAbs(src, alpha=a, beta=0)
|
|
|
|
# This code checks if there are seperate colour channels. In the past the
|
|
# macbeth locator ran on jpgs and this makes it robust to different
|
|
# filetypes. Note that running it on a jpg has 4x the pixels of the
|
|
# average bayer channel so coordinates must be doubled.
|
|
|
|
# This is best done in img_load.py in the get_patches method. The
|
|
# coordinates and image width, height must be divided by two if the
|
|
# macbeth locator has been run on a demosaicked image.
|
|
if len(src_norm.shape) == 3:
|
|
src_bw = cv2.cvtColor(src_norm, cv2.COLOR_BGR2GRAY)
|
|
else:
|
|
src_bw = src_norm
|
|
original_bw = src_bw.copy()
|
|
|
|
# Obtain image edges
|
|
sigma = 2
|
|
src_bw = cv2.GaussianBlur(src_bw, (0, 0), sigma)
|
|
t1, t2 = 50, 100
|
|
edges = cv2.Canny(src_bw, t1, t2)
|
|
|
|
# Dilate edges to prevent self-intersections in contours
|
|
k_size = 2
|
|
kernel = np.ones((k_size, k_size))
|
|
its = 1
|
|
edges = cv2.dilate(edges, kernel, iterations=its)
|
|
|
|
# Find contours in image
|
|
conts, _ = cv2.findContours(edges, cv2.RETR_TREE,
|
|
cv2.CHAIN_APPROX_NONE)
|
|
if len(conts) == 0:
|
|
raise MacbethError(
|
|
'\nWARNING: No macbeth chart found!'
|
|
'\nNo contours found in image\n'
|
|
'Possible problems:\n'
|
|
'- Macbeth chart is too dark or bright\n'
|
|
'- Macbeth chart is occluded\n'
|
|
)
|
|
|
|
# Find quadrilateral contours
|
|
epsilon = 0.07
|
|
conts_per = []
|
|
for i in range(len(conts)):
|
|
per = cv2.arcLength(conts[i], True)
|
|
poly = cv2.approxPolyDP(conts[i], epsilon * per, True)
|
|
if len(poly) == 4 and cv2.isContourConvex(poly):
|
|
conts_per.append((poly, per))
|
|
|
|
if len(conts_per) == 0:
|
|
raise MacbethError(
|
|
'\nWARNING: No macbeth chart found!'
|
|
'\nNo quadrilateral contours found'
|
|
'\nPossible problems:\n'
|
|
'- Macbeth chart is too dark or bright\n'
|
|
'- Macbeth chart is occluded\n'
|
|
'- Macbeth chart is out of camera plane\n'
|
|
)
|
|
|
|
# Sort contours by perimeter and get perimeters within percent of median
|
|
conts_per = sorted(conts_per, key=lambda x: x[1])
|
|
med_per = conts_per[int(len(conts_per) / 2)][1]
|
|
side = med_per / 4
|
|
perc = 0.1
|
|
med_low, med_high = med_per * (1 - perc), med_per * (1 + perc)
|
|
squares = []
|
|
for i in conts_per:
|
|
if med_low <= i[1] and med_high >= i[1]:
|
|
squares.append(i[0])
|
|
|
|
# Obtain coordinates of nomralised macbeth and squares
|
|
square_verts, mac_norm = get_square_verts(0.06)
|
|
# For each square guess, find 24 possible macbeth chart centres
|
|
mac_mids = []
|
|
squares_raw = []
|
|
for i in range(len(squares)):
|
|
square = squares[i]
|
|
squares_raw.append(square)
|
|
|
|
# Convert quads to rotated rectangles. This is required as the
|
|
# 'squares' are usually quite irregular quadrilaterls, so
|
|
# performing a transform would result in exaggerated warping and
|
|
# inaccurate macbeth chart centre placement
|
|
rect = cv2.minAreaRect(square)
|
|
square = cv2.boxPoints(rect).astype(np.float32)
|
|
|
|
# Reorder vertices to prevent 'hourglass shape'
|
|
square = sorted(square, key=lambda x: x[0])
|
|
square_1 = sorted(square[:2], key=lambda x: x[1])
|
|
square_2 = sorted(square[2:], key=lambda x: -x[1])
|
|
square = np.array(np.concatenate((square_1, square_2)), np.float32)
|
|
square = np.reshape(square, (4, 2)).astype(np.float32)
|
|
squares[i] = square
|
|
|
|
# Find 24 possible macbeth chart centres by trasnforming normalised
|
|
# macbeth square vertices onto candidate square vertices found in image
|
|
for j in range(len(square_verts)):
|
|
verts = square_verts[j]
|
|
p_mat = cv2.getPerspectiveTransform(verts, square)
|
|
mac_guess = cv2.perspectiveTransform(mac_norm, p_mat)
|
|
mac_guess = np.round(mac_guess).astype(np.int32)
|
|
|
|
mac_mid = np.mean(mac_guess, axis=1)
|
|
mac_mids.append([mac_mid, (i, j)])
|
|
|
|
if len(mac_mids) == 0:
|
|
raise MacbethError(
|
|
'\nWARNING: No macbeth chart found!'
|
|
'\nNo possible macbeth charts found within image'
|
|
'\nPossible problems:\n'
|
|
'- Part of the macbeth chart is outside the image\n'
|
|
'- Quadrilaterals in image background\n'
|
|
)
|
|
|
|
# Reshape data
|
|
for i in range(len(mac_mids)):
|
|
mac_mids[i][0] = mac_mids[i][0][0]
|
|
|
|
# Find where midpoints cluster to identify most likely macbeth centres
|
|
clustering = cluster.AgglomerativeClustering(
|
|
n_clusters=None,
|
|
compute_full_tree=True,
|
|
distance_threshold=side * 2
|
|
)
|
|
mac_mids_list = [x[0] for x in mac_mids]
|
|
|
|
if len(mac_mids_list) == 1:
|
|
# Special case of only one valid centre found (probably not needed)
|
|
clus_list = []
|
|
clus_list.append([mac_mids, len(mac_mids)])
|
|
|
|
else:
|
|
clustering.fit(mac_mids_list)
|
|
|
|
# Create list of all clusters
|
|
clus_list = []
|
|
if clustering.n_clusters_ > 1:
|
|
for i in range(clustering.labels_.max() + 1):
|
|
indices = [j for j, x in enumerate(clustering.labels_) if x == i]
|
|
clus = []
|
|
for index in indices:
|
|
clus.append(mac_mids[index])
|
|
clus_list.append([clus, len(clus)])
|
|
clus_list.sort(key=lambda x: -x[1])
|
|
|
|
elif clustering.n_clusters_ == 1:
|
|
# Special case of only one cluster found
|
|
clus_list.append([mac_mids, len(mac_mids)])
|
|
else:
|
|
raise MacbethError(
|
|
'\nWARNING: No macebth chart found!'
|
|
'\nNo clusters found'
|
|
'\nPossible problems:\n'
|
|
'- NA\n'
|
|
)
|
|
|
|
# Keep only clusters with enough votes
|
|
clus_len_max = clus_list[0][1]
|
|
clus_tol = 0.7
|
|
for i in range(len(clus_list)):
|
|
if clus_list[i][1] < clus_len_max * clus_tol:
|
|
clus_list = clus_list[:i]
|
|
break
|
|
cent = np.mean(clus_list[i][0], axis=0)[0]
|
|
clus_list[i].append(cent)
|
|
|
|
# Get centres of each normalised square
|
|
reference = get_square_centres(0.06)
|
|
|
|
# For each possible macbeth chart, transform image into
|
|
# normalised space and find correlation with reference
|
|
max_cor = 0
|
|
best_map = None
|
|
best_fit = None
|
|
best_cen_fit = None
|
|
best_ref_mat = None
|
|
|
|
for clus in clus_list:
|
|
clus = clus[0]
|
|
sq_cents = []
|
|
ref_cents = []
|
|
i_list = [p[1][0] for p in clus]
|
|
for point in clus:
|
|
i, j = point[1]
|
|
|
|
# Remove any square that voted for two different points within
|
|
# the same cluster. This causes the same point in the image to be
|
|
# mapped to two different reference square centres, resulting in
|
|
# a very distorted perspective transform since cv2.findHomography
|
|
# simply minimises error.
|
|
# This phenomenon is not particularly likely to occur due to the
|
|
# enforced distance threshold in the clustering fit but it is
|
|
# best to keep this in just in case.
|
|
if i_list.count(i) == 1:
|
|
square = squares_raw[i]
|
|
sq_cent = np.mean(square, axis=0)
|
|
ref_cent = reference[j]
|
|
sq_cents.append(sq_cent)
|
|
ref_cents.append(ref_cent)
|
|
|
|
# At least four squares need to have voted for a centre in
|
|
# order for a transform to be found
|
|
if len(sq_cents) < 4:
|
|
raise MacbethError(
|
|
'\nWARNING: No macbeth chart found!'
|
|
'\nNot enough squares found'
|
|
'\nPossible problems:\n'
|
|
'- Macbeth chart is occluded\n'
|
|
'- Macbeth chart is too dark of bright\n'
|
|
)
|
|
|
|
ref_cents = np.array(ref_cents)
|
|
sq_cents = np.array(sq_cents)
|
|
|
|
# Find best fit transform from normalised centres to image
|
|
h_mat, mask = cv2.findHomography(ref_cents, sq_cents)
|
|
if 'None' in str(type(h_mat)):
|
|
raise MacbethError(
|
|
'\nERROR\n'
|
|
)
|
|
|
|
# Transform normalised corners and centres into image space
|
|
mac_fit = cv2.perspectiveTransform(mac_norm, h_mat)
|
|
mac_cen_fit = cv2.perspectiveTransform(np.array([reference]), h_mat)
|
|
|
|
# Transform located corners into reference space
|
|
ref_mat = cv2.getPerspectiveTransform(
|
|
mac_fit,
|
|
np.array([ref_corns])
|
|
)
|
|
map_to_ref = cv2.warpPerspective(
|
|
original_bw, ref_mat,
|
|
(ref_w, ref_h)
|
|
)
|
|
|
|
# Normalise brigthness
|
|
a = 125 / np.average(map_to_ref)
|
|
map_to_ref = cv2.convertScaleAbs(map_to_ref, alpha=a, beta=0)
|
|
|
|
# Find correlation with bw reference macbeth
|
|
cor = correlate(map_to_ref, ref)
|
|
|
|
# Keep only if best correlation
|
|
if cor > max_cor:
|
|
max_cor = cor
|
|
best_map = map_to_ref
|
|
best_fit = mac_fit
|
|
best_cen_fit = mac_cen_fit
|
|
best_ref_mat = ref_mat
|
|
|
|
# Rotate macbeth by pi and recorrelate in case macbeth chart is
|
|
# upside-down
|
|
mac_fit_inv = np.array(
|
|
([[mac_fit[0][2], mac_fit[0][3],
|
|
mac_fit[0][0], mac_fit[0][1]]])
|
|
)
|
|
mac_cen_fit_inv = np.flip(mac_cen_fit, axis=1)
|
|
ref_mat = cv2.getPerspectiveTransform(
|
|
mac_fit_inv,
|
|
np.array([ref_corns])
|
|
)
|
|
map_to_ref = cv2.warpPerspective(
|
|
original_bw, ref_mat,
|
|
(ref_w, ref_h)
|
|
)
|
|
a = 125 / np.average(map_to_ref)
|
|
map_to_ref = cv2.convertScaleAbs(map_to_ref, alpha=a, beta=0)
|
|
cor = correlate(map_to_ref, ref)
|
|
if cor > max_cor:
|
|
max_cor = cor
|
|
best_map = map_to_ref
|
|
best_fit = mac_fit_inv
|
|
best_cen_fit = mac_cen_fit_inv
|
|
best_ref_mat = ref_mat
|
|
|
|
# Check best match is above threshold
|
|
cor_thresh = 0.6
|
|
if max_cor < cor_thresh:
|
|
raise MacbethError(
|
|
'\nWARNING: Correlation too low'
|
|
'\nPossible problems:\n'
|
|
'- Bad lighting conditions\n'
|
|
'- Macbeth chart is occluded\n'
|
|
'- Background is too noisy\n'
|
|
'- Macbeth chart is out of camera plane\n'
|
|
)
|
|
|
|
# Represent coloured macbeth in reference space
|
|
best_map_col = cv2.warpPerspective(
|
|
original, best_ref_mat, (ref_w, ref_h)
|
|
)
|
|
best_map_col = cv2.resize(
|
|
best_map_col, None, fx=4, fy=4
|
|
)
|
|
a = 125 / np.average(best_map_col)
|
|
best_map_col_norm = cv2.convertScaleAbs(
|
|
best_map_col, alpha=a, beta=0
|
|
)
|
|
|
|
# Rescale coordinates to original image size
|
|
fit_coords = (best_fit / factor, best_cen_fit / factor)
|
|
|
|
return (max_cor, best_map_col_norm, fit_coords, True)
|
|
|
|
# Catch macbeth errors and continue with code
|
|
except MacbethError as error:
|
|
# \todo: This happens so many times in a normal run, that it shadows
|
|
# all the relevant output
|
|
# logger.warning(error)
|
|
return (0, None, None, False)
|
|
|
|
|
|
def find_macbeth(img, mac_config):
|
|
small_chart = mac_config['small']
|
|
show = mac_config['show']
|
|
|
|
# Catch the warnings
|
|
warnings.simplefilter("ignore")
|
|
warnings.warn("runtime", RuntimeWarning)
|
|
|
|
# Reference macbeth chart is created that will be correlated with the
|
|
# located macbeth chart guess to produce a confidence value for the match.
|
|
script_dir = Path(os.path.realpath(os.path.dirname(__file__)))
|
|
macbeth_ref_path = script_dir.joinpath('macbeth_ref.pgm')
|
|
ref = cv2.imread(str(macbeth_ref_path), flags=cv2.IMREAD_GRAYSCALE)
|
|
ref_w = 120
|
|
ref_h = 80
|
|
rc1 = (0, 0)
|
|
rc2 = (0, ref_h)
|
|
rc3 = (ref_w, ref_h)
|
|
rc4 = (ref_w, 0)
|
|
ref_corns = np.array((rc1, rc2, rc3, rc4), np.float32)
|
|
ref_data = (ref, ref_w, ref_h, ref_corns)
|
|
|
|
# Locate macbeth chart
|
|
cor, mac, coords, ret = get_macbeth_chart(img, ref_data)
|
|
|
|
# Following bits of code try to fix common problems with simple techniques.
|
|
# If now or at any point the best correlation is of above 0.75, then
|
|
# nothing more is tried as this is a high enough confidence to ensure
|
|
# reliable macbeth square centre placement.
|
|
|
|
# Keep a list that will include this and any brightened up versions of
|
|
# the image for reuse.
|
|
all_images = [img]
|
|
|
|
for brightness in [2, 4]:
|
|
if cor >= 0.75:
|
|
break
|
|
img_br = cv2.convertScaleAbs(img, alpha=brightness, beta=0)
|
|
all_images.append(img_br)
|
|
cor_b, mac_b, coords_b, ret_b = get_macbeth_chart(img_br, ref_data)
|
|
if cor_b > cor:
|
|
cor, mac, coords, ret = cor_b, mac_b, coords_b, ret_b
|
|
|
|
# In case macbeth chart is too small, take a selection of the image and
|
|
# attempt to locate macbeth chart within that. The scale increment is
|
|
# root 2
|
|
|
|
# These variables will be used to transform the found coordinates at
|
|
# smaller scales back into the original. If ii is still -1 after this
|
|
# section that means it was not successful
|
|
ii = -1
|
|
w_best = 0
|
|
h_best = 0
|
|
d_best = 100
|
|
|
|
# d_best records the scale of the best match. Macbeth charts are only looked
|
|
# for at one scale increment smaller than the current best match in order to avoid
|
|
# unecessarily searching for macbeth charts at small scales.
|
|
# If a macbeth chart ha already been found then set d_best to 0
|
|
if cor != 0:
|
|
d_best = 0
|
|
|
|
for index, pair in enumerate([{'sel': 2 / 3, 'inc': 1 / 6},
|
|
{'sel': 1 / 2, 'inc': 1 / 8},
|
|
{'sel': 1 / 3, 'inc': 1 / 12},
|
|
{'sel': 1 / 4, 'inc': 1 / 16}]):
|
|
if cor >= 0.75:
|
|
break
|
|
|
|
# Check if we need to check macbeth charts at even smaller scales. This
|
|
# slows the code down significantly and has therefore been omitted by
|
|
# default, however it is not unusably slow so might be useful if the
|
|
# macbeth chart is too small to be picked up to by the current
|
|
# subselections. Use this for macbeth charts with side lengths around
|
|
# 1/5 image dimensions (and smaller...?) it is, however, recommended
|
|
# that macbeth charts take up as large as possible a proportion of the
|
|
# image.
|
|
if index >= 2 and (not small_chart or d_best <= index - 1):
|
|
break
|
|
|
|
w, h = list(img.shape[:2])
|
|
# Set dimensions of the subselection and the step along each axis
|
|
# between selections
|
|
w_sel = int(w * pair['sel'])
|
|
h_sel = int(h * pair['sel'])
|
|
w_inc = int(w * pair['inc'])
|
|
h_inc = int(h * pair['inc'])
|
|
|
|
loop = int(((1 - pair['sel']) / pair['inc']) + 1)
|
|
# For each subselection, look for a macbeth chart
|
|
for img_br in all_images:
|
|
for i in range(loop):
|
|
for j in range(loop):
|
|
w_s, h_s = i * w_inc, j * h_inc
|
|
img_sel = img_br[w_s:w_s + w_sel, h_s:h_s + h_sel]
|
|
cor_ij, mac_ij, coords_ij, ret_ij = get_macbeth_chart(img_sel, ref_data)
|
|
|
|
# If the correlation is better than the best then record the
|
|
# scale and current subselection at which macbeth chart was
|
|
# found. Also record the coordinates, macbeth chart and message.
|
|
if cor_ij > cor:
|
|
cor = cor_ij
|
|
mac, coords, ret = mac_ij, coords_ij, ret_ij
|
|
ii, jj = i, j
|
|
w_best, h_best = w_inc, h_inc
|
|
d_best = index + 1
|
|
|
|
# Transform coordinates from subselection to original image
|
|
if ii != -1:
|
|
for a in range(len(coords)):
|
|
for b in range(len(coords[a][0])):
|
|
coords[a][0][b][1] += ii * w_best
|
|
coords[a][0][b][0] += jj * h_best
|
|
|
|
if not ret:
|
|
return None
|
|
|
|
coords_fit = coords
|
|
if cor < 0.75:
|
|
logger.warning(f'Low confidence {cor:.3f} for macbeth chart')
|
|
|
|
if show:
|
|
draw_macbeth_results(img, coords_fit)
|
|
|
|
return coords_fit
|
|
|
|
|
|
def locate_macbeth(image: Image, config: dict):
|
|
# Find macbeth centres
|
|
av_chan = (np.mean(np.array(image.channels), axis=0) / (2**16))
|
|
av_val = np.mean(av_chan)
|
|
if av_val < image.blacklevel_16 / (2**16) + 1 / 64:
|
|
logger.warning(f'Image {image.path.name} too dark')
|
|
return None
|
|
|
|
macbeth = find_macbeth(av_chan, config['general']['macbeth'])
|
|
|
|
if macbeth is None:
|
|
logger.warning(f'No macbeth chart found in {image.path.name}')
|
|
return None
|
|
|
|
mac_cen_coords = macbeth[1]
|
|
if not image.get_patches(mac_cen_coords):
|
|
logger.warning(f'Macbeth patches have saturated in {image.path.name}')
|
|
return None
|
|
|
|
image.macbeth = macbeth
|
|
|
|
return macbeth
|