#!/usr/bin/env python
# coding: utf-8

# In[1]:


import cv2 as cv
import numpy as np
from matplotlib import pyplot as plt
import os
from tqdm import tqdm


# In[2]:


source_dir = '/home/zoe/photos/2022/12_germering'
target_dir = os.path.join(source_dir, 'pythonomorphic')
os.makedirs(target_dir, exist_ok=True)


# In[ ]:


for fn in tqdm(os.listdir(source_dir), unit='image'):
    if fn.lower().endswith('.jpg'):
        infile = os.path.join(source_dir, fn)
        outfile = os.path.join(target_dir, fn)

        if os.path.isfile(outfile):
            continue

        img = cv.imread(infile)
        imsize = np.min(img.shape[:2])

        img_bw = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
        img_blurred = cv.medianBlur(img_bw, 2 * int(imsize / 200) + 1)
        threshold = int(np.max(img_blurred) * 0.99)
        ret, th1 = cv.threshold(img_blurred, threshold, 255, cv.THRESH_BINARY)

        kernel = np.ones((int(imsize/200), int(imsize/200)), np.uint8)

        blurred_highlights = cv.equalizeHist(
                                    cv.GaussianBlur(
                                        cv.morphologyEx(
                                            th1, cv.MORPH_OPEN, kernel),
                                        (int(imsize/7.4) * 2 + 1, int(imsize/320) * 2 + 1), 0))

        th1_blur_grey = cv.GaussianBlur(blurred_highlights,
                                        (int(imsize/200) * 2 + 1, int(imsize/200) * 2 + 1), 0)

        colored_flares = cv.applyColorMap(th1_blur_grey, cv.COLORMAP_OCEAN)
        dst = cv.addWeighted(img, 1, colored_flares, 0.5, 0)

        cv.imwrite(outfile, dst, [cv.IMWRITE_JPEG_QUALITY, 96])
        # print('.', end='', flush=True)