# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# ## Make Colorgram Continuum Notebook
#
# This notebook will generate a colorgram continuum using a sliding window across an image set.

# +
import compsyn
from PIL import Image
import os

import PIL
import numpy as np
from numba import jit, prange
from collections import OrderedDict 
# -

from compsyn.helperfunctions import settings, search_and_download, run_google_vision, write_img_classifications_to_file

from compsyn.wordnet_functions import get_wordnet_tree_data

GOOGLE_APPLICATION_CREDENTIALS = "compsyn3-8cf6580619a9.json"

DRIVER_PATH = "chromedriver"

settings(GOOGLE_APPLICATION_CREDENTIALS, DRIVER_PATH)

number_images = 100 
home = os.getcwd()

search_terms = ['tiles']

img_urls_dict = {}
for search_term in search_terms:
    print(search_term)
    urls = search_and_download(search_term = search_term, driver_path = DRIVER_PATH, home = home, number_images = number_images)
    img_urls_dict[search_term] = urls


# ### Create new continuum folder 

def make_continuum(term1, term2, path1, path2): 
    newpath = term1 + '_' + term2 + '_continuum'
    continuum_path = 'C:\\Users\\Doug\\Desktop\\comp-syn\\downloads\\' + newpath
    
    if not os.path.isdir(continuum_path): 
        os.mkdir(continuum_path)
    
    imgset1 = os.listdir(path1)
    imgset2 = os.listdir(path2)
    
    i=0
    os.chdir(path1)
    for img in imgset1: 
        img = Image.open(img) 
        img.save(continuum_path + '\\' + "img_" + str(i) + '.jpg', 'JPEG')
        i+=1
        
    os.chdir(path2)
    for img in imgset2: 
        img = Image.open(img) 
        img.save(continuum_path + '\\' + "img_" + str(i) + '.jpg', 'JPEG')
        i+=1
    
    return continuum_path


# +
term1 = 'morning ocean'
term2 = 'night ocean'

path1 = 'C:\\Users\\Doug\\Desktop\\comp-syn\\downloads\\' + term1
path2 = 'C:\\Users\\Doug\\Desktop\\comp-syn\\downloads\\' + term2
# -

continuum_path = make_continuum(term1, term2, path1, path2)

continuum_path = 'C:\\Users\\Doug\\Desktop\\comp-syn\\downloads\\dead tree_tree_continuum'

# ### Load new continuum folder via load_img_dict_with_window

# +
window=60
num_imgs = len(os.listdir(continuum_path))
idxs = range(0, len(os.listdir(continuum_path)) - window)

savepath = 'C:\\Users\\Doug\\Desktop\\comp-syn\\colorgrams\\'
savepath = savepath + term1 + "_" + term2 + "_CGs_" + str(window)

if not os.path.isdir(savepath): 
    os.mkdir(savepath)

# +
#Arrange Continuum Images so they are loaded in the proper order 
continuum_files = os.listdir(continuum_path)

numIDs = list(map(lambda sub:int(''.join([ele for ele in sub if ele.isnumeric()])), continuum_files)) 

continuum_dict={}
for i in range(len(numIDs)):
    ID = numIDs[i]
    continuum_dict[ID] = continuum_files[i]

continuum_dict_ord = dict(OrderedDict(sorted(continuum_dict.items())))
continuum_files = list(continuum_dict_ord.values())

# -

@jit
def make_window_color_gram(idx, continuum_path = continuum_path, 
                           continuum_files = continuum_files, 
                           window=window, savepath=savepath): 
    
    imagedata.load_image_continuum_from_folder(continuum_path, continuum_files, idx, window)
    imageanalysis = compsyn.analysis.ImageAnalysis(imagedata)    
    imageanalysis.compute_color_distributions(imageanalysis.labels_list, color_rep=['jzazbz', 'rgb'])
    compressed_image_dict = imageanalysis.get_composite_image()
    
    img = PIL.Image.fromarray(compressed_image_dict[list(compressed_image_dict.keys())[0]].astype(np.uint8))
    img.save(savepath + '\\' + "CG_" + str(idx) + '.jpg', 'JPEG')



imagedata = compsyn.datahelper.ImageData()

for idx in idxs: 
    make_window_color_gram(idx, continuum_path = continuum_path, 
                           continuum_files = continuum_files,
                           window=window, savepath=savepath)

StarcoderdataJupyterScriptsDedupFiltered