Source code for tcutility.report.figure_resizer

'''
This module is used for resizing pictures containing molecules.
'''
import numpy as np
import cv2
import matplotlib.pyplot as plt
import os
from typing import Optional, Dict, Union


def _analyse_img(file, plot=False):
    '''
    Function used for analysing and getting useful information from an image.
    This includes circle locations and sizes.
    '''
    # Read image
    img = cv2.imread(file, cv2.IMREAD_UNCHANGED)
    img = _remove_padding(img)
    # Smooth it
    img_copy = img.copy()

    # Convert to greyscale
    img_gray = cv2.medianBlur(img_copy, 3)
    img_gray = cv2.cvtColor(img_gray, cv2.COLOR_BGRA2GRAY)

    # Apply Hough transform to greyscale image
    circles = cv2.HoughCircles(img_gray, cv2.HOUGH_GRADIENT_ALT, 1.5, 100,
                               param1=400, param2=0.8, minRadius=0, maxRadius=800)
    circles = np.uint16(np.around(circles))[0]
    circles = circles[(-circles[:, 2]).argsort()]

    if plot:
        img_copy = cv2.cvtColor(img_copy, cv2.COLOR_BGR2RGB)
        for i, circle in enumerate(circles):
            # draw the outer circle
            cv2.circle(img_copy,(circle[0],circle[1]),circle[2],(0,255,0),2)
            # draw the center of the circle
            cv2.circle(img_copy,(circle[0],circle[1]),2,(0,0,255),3)
            cv2.putText(img_copy, str(i), (circle[0] - 40, circle[1] + 44), cv2.FONT_HERSHEY_SIMPLEX, 4, (0, 0, 0), 16, cv2.LINE_AA)
            cv2.putText(img_copy, str(i), (circle[0] - 40, circle[1] + 44), cv2.FONT_HERSHEY_SIMPLEX, 4, (255, 255, 255), 8, cv2.LINE_AA)

        plt.imshow(img_copy)
        plt.draw()
        plt.pause(0.001)
    return circles, img


def _remove_padding(img):
    '''
    Function used to remove padding from an image.
    '''
    gray = cv2.cvtColor(img, cv2.COLOR_BGRA2GRAY)
    gray = 255*(gray < 240).astype(np.uint8) # To invert the text to white
    coords = cv2.findNonZero(gray) # Find all non-zero points (text)
    x, y, w, h = cv2.boundingRect(coords) # Find minimum spanning bounding box
    rect = img[y:y+h, x:x+w] # Crop the image - note we do this on the original image
    return rect


[docs] def resize(img_paths, circle_numbers: Optional[Dict] = None, padding: Union[str, int, float] = 0): ''' The main function for this module. Takes a directory `d` and selected circles and resizes and moves images in order to produce new aligned images. Args: d: the directory to take images from. circle_numbers: a dictionary containing the image names as keys and circle indices as the values. The new images will be aligned based on these selected circles. padding: the amount of padding to add to the new images. If given an integer, it adds `padding` number of pixels. Alternatively a string consisting of a number and a `%` sign can be given to add relative padding. E.g. `padding='10%'` will add a 10% padding to the images. ''' circles = {} imgs = {} for file in img_paths: if file not in circle_numbers: continue circles_, img_ = _analyse_img(file) circles[file] = circles_[circle_numbers[file]] imgs[file] = img_ # get the atom sizes, which is simply the radius of the circle we selected atom_sizes = {file: circle[2] for file, circle in circles.items()} # calculate ratios for resizing the images ratios = {file: max(atom_sizes.values())/size for file, size in atom_sizes.items()} # resize old images to make new ones imgs = {file: cv2.resize(img, [int(round(img.shape[1] * ratios[file])), int(round(img.shape[0] * ratios[file]))]) for file, img in imgs.items()} # also resize circles to match the new ratio circles = {file: circle * ratios[file] for file, circle in circles.items()} # now find out where to place the circles circle_x = [circle[0] for circle in circles.values()] circle_y = [circle[1] for file, circle in circles.items()] pad_left = {file: int((max(circle_x) - circle[0])) for file, circle in circles.items()} pad_top = {file: int((max(circle_y) - circle[1])) for file, circle in circles.items()} pad = {file: [(pad_top[file], 0), (pad_left[file], 0), (0, 0)] for file in imgs} imgs = {file: np.pad(img, pad[file], constant_values=0) for file, img in imgs.items()} imsize_x = [img.shape[1] for img in imgs.values()] imsize_y = [img.shape[0] for img in imgs.values()] pad_right = {file: int(max(imsize_x) - img.shape[1]) for file, img in imgs.items()} pad_bottom = {file: int(max(imsize_y) - img.shape[0]) for file, img in imgs.items()} pad = {file: [(0, pad_bottom[file]), (0, pad_right[file]), (0, 0)] for file in imgs} imgs = {file: np.pad(img, pad[file], constant_values=0) for file, img in imgs.items()} # add final padding if isinstance(padding, str) and padding.endswith('%'): padding = float(padding.removesuffix('%'))/100 padding = int(padding * list(imgs.values())[0].shape[1]), int(padding * list(imgs.values())[0].shape[0]) else: padding = int(padding), int(padding) pad = {file: [(padding[0], padding[0]), (padding[1], padding[1]), (0, 0)] for file in imgs} imgs = {file: np.pad(img, pad[file], constant_values=0) for file, img in imgs.items()} ret = {} for file, img in imgs.items(): new_file = os.path.join(os.path.split(file)[0], 'resized_' + os.path.split(file)[1]) cv2.imwrite(new_file, img) ret[file] = new_file return ret