Last active
December 14, 2022 15:56
-
-
Save Stovoy/bf0f792c653b2767e0f40b3380c0fdb6 to your computer and use it in GitHub Desktop.
A script to correct the rotation of pennies using OpenCV
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import argparse | |
import math | |
import multiprocessing | |
import os | |
import itertools as it | |
from scipy import ndimage | |
from skimage.metrics import structural_similarity | |
import cv2 | |
import numpy as np | |
import matplotlib.pyplot as plt | |
class PennyRotator: | |
def __init__(self): | |
self.face_template = cv2.imread('face_template.png') | |
self.face_template = self.normalize(self.face_template) | |
def correct_rotation_on_penny(self, i, count, image_path): | |
print(f'[{i + 1} / {count}] Correcting rotation for "{image_path}"') | |
# Open the image | |
original_image = cv2.imread(image_path) | |
normalized_image = self.normalize(original_image) | |
print(f'[{i + 1}] Starting rotate and compare') | |
# Get all angles between 0 and 360 inclusive with an increment of 1 into an array | |
angles = np.arange(1, 360, 1) | |
best_angle = 0 | |
best_score = 0 | |
for angle in angles: | |
angle, score = self.rotate_and_compare(normalized_image, angle) | |
if score > best_score: | |
best_score = score | |
best_angle = angle | |
print(f'[{i + 1}] Best Angle: {best_angle}, Best Score: {best_score}') | |
result_image = self.rotate(original_image, best_angle) | |
return result_image, original_image, normalized_image, image_path | |
def center(self, image): | |
return tuple(np.array(image.shape[1::-1]) / 2) | |
def rotate_and_compare(self, image, angle): | |
rotated = self.rotate(image, angle) | |
score = structural_similarity(self.face_template, rotated, gaussian_weights=True) | |
return angle, score | |
def rotate(self, image, angle): | |
rotation_matrix = cv2.getRotationMatrix2D(self.center(image), angle, 1.0) | |
return cv2.warpAffine(image, rotation_matrix, image.shape[1::-1], flags=cv2.INTER_LINEAR) | |
def normalize(self, image): | |
# Apply CLAHE technique to normalize the image | |
image_lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB) | |
l_channel, a, b = cv2.split(image_lab) | |
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)) | |
l_channel = clahe.apply(l_channel) | |
image_lab = cv2.merge((l_channel, a, b)) | |
image = cv2.cvtColor(cv2.cvtColor(image_lab, cv2.COLOR_LAB2BGR), cv2.COLOR_BGR2GRAY) | |
size = 100 | |
image = cv2.resize(image, (size, size), interpolation=cv2.INTER_AREA) | |
# Try Canny approach, but use more and more median filtering while it's too noisy. | |
retries_with_filtering = 5 | |
while True: | |
result = cv2.Canny(image, 90, 150, apertureSize=3) | |
brightness = cv2.sumElems(result)[0] | |
if brightness <= 400_000 or retries_with_filtering <= 0: | |
break | |
image = ndimage.median_filter(image, 3) | |
retries_with_filtering -= 1 | |
image = result | |
# Mask the inner circle | |
mask = cv2.circle(np.zeros_like(image), (size // 2, size // 2), int(size // 2.5), (255, 255, 255), -1) | |
image = cv2.bitwise_and(image, mask) | |
return image | |
parser = argparse.ArgumentParser() | |
parser.add_argument("-p", "--plot", help="Show plots", action="store_true") | |
parser.set_defaults(plot=False) | |
args = parser.parse_args() | |
def main(): | |
rotator = PennyRotator() | |
if not os.path.exists('output'): | |
os.makedirs('output') | |
files = os.listdir('data') | |
with multiprocessing.Pool() as pool: | |
async_results = [] | |
for i, file in enumerate(files): | |
async_results.append( | |
pool.apply_async(rotator.correct_rotation_on_penny, (i, len(files), os.path.join('data', file),))) | |
for i, result in enumerate(async_results): | |
print(f'[{i + 1} / {len(async_results)}] Processing result ') | |
result_image, original_image, normalized_image, image_path = result.get() | |
image_path = os.path.basename(image_path) | |
cv2.imwrite(os.path.join('output', image_path), result_image) | |
if args.plot: | |
_, axes = plt.subplots(2, 2, figsize=(8, 8)) | |
axes = axes.flatten() | |
titles = ['Original', 'Result', 'Original Normalized', 'Template Normalized'] | |
is_gray = [False, False, True, True] | |
images = [ | |
original_image, | |
result_image, | |
normalized_image, | |
rotator.face_template | |
] | |
for axis, image, title, is_gray in zip(axes, images, titles, is_gray): | |
axis.set_title(title) | |
if is_gray: | |
axis.imshow(image, cmap='gray') | |
else: | |
axis.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) | |
plt.suptitle(image_path) | |
plt.show() | |
if __name__ == '__main__': | |
main() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Template image