<jupyter_start><jupyter_text><jupyter_code>!nvidia-smi -L
from google.colab import drive
drive.mount('/content/drive')
%%writefile requirements.txt
numpy>=1.7.0
imutils==0.5.2
opencv-python>=4.1.0.25
scipy
lpips
torch==1.7.1
torchvision==0.8.2
!pip install -r requirements.txt
%%writefile lpips-improve-scale.py

import argparse
import queue
import lpips
import torch
import os
import cv2
import numpy as np
import _thread

def saveImage(img,path,filename):
	if(args.file_extension == "png"):
		new_file = os.path.splitext(filename)[0] + ".png"
		cv2.imwrite(os.path.join(path, new_file), img, [cv2.IMWRITE_PNG_COMPRESSION, 0])
	elif(args.file_extension == "jpg"):
		new_file = os.path.splitext(filename)[0] + ".jpg"
		cv2.imwrite(os.path.join(path, new_file), img, [cv2.IMWRITE_JPEG_QUALITY, 90])

def parse_args():
    desc = "Tools to crop unnecessary space from outside of images" 
    parser = argparse.ArgumentParser(description=desc)

    parser.add_argument('-i','--input_folder', type=str,
        default='./input/',
        help='Directory path to the inputs folder. (default: %(default)s)')

    parser.add_argument('-c', '--compare_folder', type=str,
        default='./compare/',
        help='Directory path to the compare folder. (default: %(default)s)')

    parser.add_argument('-o','--output_folder', type=str,
        default='./output/',
        help='Directory path to the outputs folder. (default: %(default)s)')

    parser.add_argument('--file_extension', type=str,
        default='png',
        help='file extension ["png","jpg"] (default: %(default)s)')

    parser.add_argument('--verbose', action='store_true',
        help='Print progress to console.')

    parser.add_argument('--use_gpu', action='store_true',
        help='Use GPU for lpips.')

    parser.add_argument('-j' '--jobs', type=int,
        default=1,
        help='The number of threads to use. (default: %(default)s)')

    args = parser.parse_args()
    return args

def threadRunner(threadName):
    while(not q.empty()):
        filename = q.get()
        optimize_image(threadName, filename)

def optimize_image(threadName, filename):
    file_path = os.path.join(rootDir, filename)
    if(args.verbose): print('(%s) processing\t- file %s (full path: %s)' % (threadName, filename, file_path))
            
    img = lpips.load_image(file_path)
    # print(type(img))

    r,g,b = cv2.split(img)
    img_bgr = cv2.merge([b,g,r])

    scaled_images = []
    for i in np.arange(0.8, 1.3, 0.1):
        scaled_images.append(scale_image(img_bgr, i))
        # saveImage(scaled_images[-1],args.output_folder,str(i) + filename)


    # for i in range(len(scaled_images)):

    tensors = []
    for i in range(len(scaled_images)):
        b, g, r = cv2.split(scaled_images[i])
        tensor = lpips.im2tensor(cv2.merge([r, g, b]))
        if args.use_gpu:
            tensor = tensor.cuda()
        tensors.append(tensor)

    values = []
    for i in range(len(tensors)):
        val = 0.0
        for j in range(len(compare_imgs)):
            # print(tensors[i], compare_imgs[j])

            img0 = tensors[i]
            img1 = compare_imgs[j]

            # if(args.use_gpu):
            #     img0 = img0.cuda()
            #     img1 = img1.cuda()

            val += loss_fn_vgg.forward(img0, img1).detach().item()


        values.append(val/len(compare_imgs))

    maxIndex = 0
    maxVal = 0.0
    for i in range(len(values)):
        if (values[i] > maxVal):
            maxIndex = i
            maxVal = values[i]

    if args.verbose:
        print(values)
        print('saving index %d' % maxIndex)

    saveImage(scaled_images[maxIndex],args.output_folder,filename)




    # if hasattr(img, 'copy'):
    #     processImage(img,os.path.splitext(filename)[0])

def scale_image(image: np.ndarray, scale: float) -> np.ndarray:
        if scale == 0:
            return image.copy()
        out = image.copy()
        start_shape = out.shape
        dim = (int(out.shape[1] * scale), int(out.shape[0] * scale))
        if scale > 1:
            out = cv2.resize(out, dim, interpolation=cv2.INTER_AREA)
            x_crop = (dim[0] - start_shape[1]) // 2
            y_crop = (dim[1] - start_shape[0]) // 2
            out = out[y_crop:dim[1] - y_crop, x_crop:dim[0] - x_crop]
        else:
            out = cv2.resize(out, dim, interpolation=cv2.INTER_AREA)

            pad_x = (start_shape[1] - dim[0]) // 2
            pad_y = (start_shape[0] - dim[1]) // 2

            in_img_h, in_img_w, channels = out.shape
            out_img_h = in_img_h + (pad_y * 2)
            out_img_w = in_img_w + (pad_x * 2)
            mask = np.zeros((out_img_h, out_img_w, 1), np.uint8)
            fill_color = (255, 255, 255)
            out_img = cv2.copyMakeBorder(out, pad_y, pad_y, pad_x, pad_x, cv2.BORDER_CONSTANT, value=fill_color)
            mask = cv2.rectangle(mask, (0, 0), (pad_x - 1, out_img_h - 1), (255), -1)
            mask = cv2.rectangle(mask, (0, 0), (out_img_w - 1, pad_y - 1), (255), -1)
            mask = cv2.rectangle(mask, (0, out_img_h - pad_y), (out_img_w - 1, out_img_h - 1), (255), -1)
            mask = cv2.rectangle(mask, (out_img_w - pad_x, 0), (out_img_w - 1, out_img_h - 1), (255), -1)

            out = cv2.inpaint(out_img, mask, 3, cv2.INPAINT_TELEA)

        if out.shape[1] != 1024 or out.shape[0] != 1024:
            out = cv2.resize(out, (1024, 1024), interpolation=cv2.INTER_AREA)

        return out

def main():
    global args
    global q
    global rootDir
    global compare_imgs
    global loss_fn_vgg
    args = parse_args()
    q = queue.Queue()
    compare_imgs = []
    loss_fn_vgg = lpips.LPIPS(net='vgg', version=0.1)

    if(args.use_gpu):
	    loss_fn_vgg.cuda()

    if os.path.isdir(args.input_folder):
        print("Processing folder: " + args.input_folder)
    else:
        print("Not a working input_folder path: " + args.input_folder)
        return

    if os.path.isdir(args.compare_folder):
        print("Compare folder: " + args.compare_folder)
    else:
        print("Not a working compare_folder path: " + args.compare_folder)
        return

    os.makedirs(args.output_folder, exist_ok=True)

    for root, subdirs, files in os.walk(args.compare_folder):
        if(args.verbose): print('--\nroot = ' + root)

        for subdir in subdirs:
            if(args.verbose): print('\t- subdirectory ' + subdir)

        for filename in files:
            # add files to queue
            file_path = os.path.join(root, filename)
            img = lpips.load_image(file_path)
            # print(img)
            tensor = lpips.im2tensor(img)
            if args.use_gpu:
                tensor = tensor.cuda()
            compare_imgs.append(tensor)
            

    for root, subdirs, files in os.walk(args.input_folder):
        if(args.verbose): print('--\nroot = ' + root)
        rootDir = root

        for subdir in subdirs:
            if(args.verbose): print('\t- subdirectory ' + subdir)

        for filename in files:
            # add files to queue
            q.put(filename)

    # start threads
    for i in range(args.j__jobs):
        try:
            threadName = 'thread-' + str(i)
            if(args.verbose): print('starting thread %s' % (threadName))
            _thread.start_new_thread(threadRunner, (threadName,))
        except:
            print("error! unable to start thread.")

    while (not q.empty()):
        pass


if __name__ == "__main__":
    main()

!python lpips-improve-scale.py --help
input_folder = '/content/drive/MyDrive/dataset-creation/treeflowers/raw-10241024'
compare_folder = '/content/drive/MyDrive/dataset-creation/treeflowers/compare-1024'
output_folder = '/content/drive/MyDrive/dataset-creation/treeflowers/out'
!python lpips-improve-scale.py -i {input_folder} -c {compare_folder} -o {output_folder} --verbose --use_gpu<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14198705748001522, 'fraction_numerical': 0.020682654485471387, 'mean_word_length': 2.5093499554764023, 'pattern_counts': {'":': 1, '<': 5, '<?xml version=': 0, '>': 10, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 1, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '560048', 'license': 'MIT', 'n_tokens_mistral': 2857, 'n_tokens_neox': 2717, 'n_words': 572, 'path': 'lpips/lpips_improve_scale.ipynb', 'repo_name': 'duskvirkus/colab-notebooks', 'size': '7881'}