[#332] Added the new tc resize CLI program

src/tcutility/cli_scripts/tcparser.py

@@ -1,4 +1,4 @@
-from tcutility.cli_scripts import read, job_script, concatenate_irc, cite, geo
+from tcutility.cli_scripts import read, job_script, concatenate_irc, cite, geo, resize_figures
 
 # to add a script:
 # 1. Add a create_subparser function and main function to your script.
@@ -10,6 +10,7 @@
     "concat-irc": concatenate_irc,
     "cite": cite,
     "geo": geo,
+    "resize": resize_figures,
 }
 
 

src/tcutility/report/figure_resizer.py

@@ -0,0 +1,135 @@
+'''
+This module is used for resizing pictures containing molecules.
+'''
+import numpy as np
+import cv2
+import matplotlib.pyplot as plt
+import os
+
+
+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, 1, 100,
+                               param1=60, param2=40, minRadius=0, maxRadius=200)
+    circles = np.uint16(np.around(circles))[0]
+    circles = circles[(-circles[:, 2]).argsort()]
+    x = np.arange(0, img.shape[1])
+    y = np.arange(0, img.shape[0])
+
+    # Draw the circles
+    colors = []
+    radii = []
+
+    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
+
+
+def resize(d, circle_numbers: dict = None, padding: str or int = 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 os.listdir(d):
+        if file == '.DS_Store':
+            continue
+        if file not in circle_numbers:
+            continue
+        circles_, img_ = _analyse_img(os.path.join(d, 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()}
+
+    os.makedirs(d + '_fixed', exist_ok=True)
+    for file, img in imgs.items():
+        cv2.imwrite(os.path.join(d + '_fixed', file), img)
+
+    cv2.imwrite(os.path.join(d + '_fixed', 'empty.png'), np.zeros_like(img) + np.array([255, 255, 255, 0]))
+    ret = {file: os.path.join(d + '_fixed', file) for file in imgs}
+    ret['empty.png'] = os.path.join(d + '_fixed', 'empty.png')
+    return ret