CoastsnapAuto/coastsnap/tag_registered.py

import datetime
import shutil
from enum import Enum
from functools import cached_property
from itertools import product
from pathlib import Path

import attr
import cv2
import numpy as np
import pandas as pd
from loguru import logger
from moviepy.editor import *

import os
os.chdir(Path(__file__).parent.resolve())

from utils import divide_chunks, datenum_to_datetime, nearest, progressbar, RegisteredImage
import time # progress bar

from PIL import Image, ImageFont, ImageDraw
from time import strptime
from dateutil import tz

import openpyxl
import scipy.io as sio
from datetime import datetime, timedelta

'''
SCRIPT LOGIC

For each site in Images Parent Directory:
    Iterate through years in reverse:
        Iterate through images names (dates) in reverse:
            Check if image name in 'Registered'? 
            If not, tag and save
            If yes, skip to the next site 
            
'''


#-----------------------------------------------------------------------------#
#          Update this file path for Image Tagging - Tide Data
# Example: images_dir = '/Users/admin/OneDrive - UNSW/My files/CoastSnap/'

code_images_dir = str(Path(os.getcwd()).parent)
sites_csv_path = os.path.join(code_images_dir, "coastsnap_sites.csv")

# RETRIEVE IMAGES PARENT DIRECTORY IN coastsnap_sites.csv

coastsnap_sites_csv = pd.read_csv(sites_csv_path)
images_parent_dir = coastsnap_sites_csv.parent_directory[0] 
images_dir = os.path.join(images_parent_dir, "Images")
#-----------------------------------------------------------------------------#   


def get_site_tide_data(images_parent_dir, site):
    
    # Retrieve tide data for the given site
    db = openpyxl.load_workbook(images_parent_dir + "/Database/CoastSnapDB.xlsx")
    beach_data = db[site]
    tide_filename = beach_data["B24"].value
    
    if tide_filename == 'NO_TIDE.mat':
        return False
        
    mat_path = images_parent_dir + '/Tide Data/' + tide_filename
    mat = sio.loadmat(mat_path)
    tide_dict = mat['tide']
    ts = tide_dict[0][0] # Format of tide/time Matlab data is gross
    tt = list(ts[0])     # ts[0] = tides, ts[1] = times, ts[2] = 
    
    print(site + ": " +"Loading tide data... (this may take half a minute)")     
    tide_times = [datenum_to_datetime(i) for i in tt] # THIS STEP TAKES A LONG TIME

    return tide_times, ts    


def tag_image(image2tag):
    
    image = Image.open(image2tag.pathname)
    draw = ImageDraw.Draw(image)
    image2tag.get_dimensions()
    
    # White Text Box 
    rect_height = image2tag.height/20
         
    # Create white text box
    draw.rectangle((0, 0, image2tag.width, rect_height), fill='white')

    # Tag image with text
    draw.text((20, rect_height/4),image2tag.tag, font = image2tag.font, fill=(0, 0, 0)) 
    
    registered_path_wrong = image2tag.pathname[:-4] + '_registered.jpg' 
    registered_path = registered_path_wrong.replace('Photoshop', 'Registered')
    image = image.convert('RGB')

    image.save(registered_path)

for site in os.listdir(images_dir): # Loop through SITES
    i=0
    site_complete = False # A flag to stop tagging images when found
    tide_data = False
    font = None
    photoshop_path = images_dir +'/'+ site + '/Photoshop'
    try: # Check if site contains 'Processed' directory
        years_list = os.listdir(photoshop_path)
        years_list.reverse()
    except: 
        continue
    for year in years_list: # Loop through YEARS
        if site_complete: break 
        year_path = photoshop_path + '/' + year
        image_list = os.listdir(year_path)
        image_list.reverse()
        for image_filename in image_list: # Loop through IMAGES
            #print(image_filename) 
            registered_year_path = year_path.replace('Photoshop', 'Registered')
            registered_image_path = registered_year_path + '/' + image_filename[:-4] + '_registered.jpg' 
            if site not in image_filename:
                continue
            # Check if image already tagged
            if os.path.isfile(registered_image_path): 
                site_complete = True;
                if i == 0:
                    print(site + ": " + str(i) + " images tagged")
                elif image2tag.tide:
                    print(site + ": " + str(i) + " images tagged with tide")
                else:
                    print(site + ": " + str(i) + " images tagged no tide")                  
                break
            else:
                if not os.path.exists(registered_year_path): # Chech that registered/year directory exists
                    os.makedirs(registered_year_path)     # if not, create it
                i += 1
                pathname = os.path.join(year_path, image_filename)
                image2tag = RegisteredImage(pathname, image_filename) # Create image object
                if i == 1: # Retrieve Tide Data once for each site
                    site_tide_data, ts = get_site_tide_data(images_parent_dir, site)

                font = image2tag.get_font()
                    
                if site_tide_data:
                    image2tag.get_tide(site_tide_data, ts)
                    
                image2tag.create_tag()
                tag_image(image2tag)
    
    if not site_complete:
        if i == 0:
            print(site + ": " + str(i) + " images tagged")
        elif image2tag.tide:
            print(site + ": " + str(i) + " images tagged with tide")
        else:
            print(site + ": " + str(i) + " images tagged no tide")