CoastsnapAuto/coastsnap/tag_registered.py

"""
test register cli
"""
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
import typer
from loguru import logger
from moviepy.editor import *

from utils import divide_chunks, datenum_to_datetime, nearest

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

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

# Enables command-line interface
app = typer.Typer()

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

parent_dir = r"C:\Users\z5079346\OneDrive - UNSW\My files\CoastSnap"

#-----------------------------------------------------------------------------#   

@app.command()
def tagging(
    folder: Path = typer.Argument(None, help="Folder with images"),
    photoshop: bool = typer.Option(True, help="Have these images been registered with photoshop?"),
    tide: bool = typer.Option(False, help="Do you want to add the tide to the image tag?")

):
    """
    Tags images based on file name.
    - Requires font file in coastsnap/fonts directory
    """

    from PIL import Image, ImageFont, ImageDraw # Throws errors if this isn't here..

    # Get image paths
    img_paths = [x for x in Path(folder).glob("*.jpg")]
    img_names = [str(x) for x in img_paths]
    logger.info(f"Tagging {len(img_names)} images")
    
    # Initialise white text box
    rect_height = 1
    width = 1

    # Check whether the directory 'tagged' exists or not
    path_name = img_names[0] 
    tagged_dir = str(Path(path_name).parent) + '/tagged'
    isExist = os.path.exists(tagged_dir)
    if not isExist:
      # Create a new directory because it does not exist 
      os.makedirs(tagged_dir)



    # # Set Image Label Info Based on Target Image
    # target_image_path = str(Path(path_name).parent.parent.parent) + r"\Target Image\Target.JPG"
    # print(target_image_path)
    # image = Image.open(target_image_path)
    # img_fraction = 0.6 # Change this to change the font size
    # fontsize = 1
    
    # # White Text Box 
    # width, height = image.size
    # rect_height = height/20
    # rect_width = width/1.5            
    # font = ImageFont.truetype("fonts/Courier New Bold.ttf", fontsize) 

    # while font.getsize(txt)[0] < img_fraction*image.size[0]:
    #     # iterate until the text size is just larger than the criteria
    #     fontsize += 1
    #     font = ImageFont.truetype("fonts/Courier New Bold.ttf", fontsize) 

    #"C:\Users\z5079346\OneDrive - UNSW\My files\CoastSnap\Images\cathieillaroo\Registered\2022"
    #"C:\Users\z5079346\OneDrive - UNSW\My files\CoastSnap\Images\cathieillaroo\Target Image\Target.JPG"



    # Tide Data
    if tide:
        # Retrieve the site name from the first image
        filename = Path(img_names[0]).name
        if photoshop == True:
            fname = filename[6:-4] 
        else:
            fname = filename
        filename_list = fname.split(".")
        site_name = filename_list[6]
        
        # Retrieve tide data for the given site
        db = openpyxl.load_workbook(parent_dir + "Database/CoastSnapDB.xlsx")
        beach_data = db[site_name]
        tide_filename = beach_data["B24"].value
        mat_path = parent_dir + 'Tide Data/' + tide_filename
        print("Loading tide data... (this may take half a minute)")
        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] = 
        tide_times = [datenum_to_datetime(i) for i in tt]
      
    fontsize = 1
    for index, img in enumerate(sorted(img_names)):
        print("Image " + str(index+1))
        filename = str(Path(img).name)
        image = Image.open(img)
        draw = ImageDraw.Draw(image)
        
        
        # Retrieve tag information from file name
        if photoshop == True:
            fname = filename[6:-4] 
        else:
            fname = filename
        
        filename_list = fname.split(".")
        posix_time = filename_list[0]
        date = filename_list[3].split("_")
        hour = date[1]
        minute = date[2]
        second = date[3]
        day = date[0]
        month = '{:02d}'.format(strptime(filename_list[2],'%b').tm_mon) # Ensure 2-digit format
        year = filename_list[5]
        timezone = filename_list[4]
        if 'snap' in filename_list:
            contributor = filename_list[8] # Mitch filename format
        else:
            contributor = filename_list[6] # Leaman filename format
        
        # Retrieve tide data 
        if tide:

            # Account for daylight savings
            # ASSUMPTION: All .mat tide files are in AEST. 
            if timezone == 'AEDT':
                hour = str(int(hour) - 1)
            date_string = year + '-' + month + '-' + day + ' ' + hour + ':' + minute + ':' + second
            img_datetime = datetime.strptime(date_string, "%Y-%m-%d %H:%M:%S") # Image date/time as a datetime object
            
            tide_date = nearest(tide_times, img_datetime)
            mat_index = tide_times.index(tide_date)  # Retrieve the index of the .mat tide/time
            mat_tide = round(ts[1][mat_index][0], 2) # Associated tide
            
            print('Image date/time: ' + date_string)
            print('Tide record:     ' + str(tide_date))
        
        
        # Image tag
        if tide:
            txt = ('Date:' + year + '/' + month + '/' + day +  
                   ' Time:' + hour + ':' + minute + 
                   ' Contributor:' + contributor + 
                   ' Tide:' + str(mat_tide) + 'm AHD')
        else:
            txt = ('Date:' + year + '/' + month + '/' + day +  
                   ' Time:' + hour + ':' + minute + 
                   ' Contributor:' + contributor)

        
        # Set the fontsize, such that the tag covers 50% the width of the first image
        if index == 0:
            img_fraction = 0.6 # Change this to change the font size
            
            # White Text Box 
            width, height = image.size
            rect_height = height/20
            rect_width = width/1.5            
            font = ImageFont.truetype("fonts/Courier New Bold.ttf", fontsize) 

            while font.getsize(txt)[0] < img_fraction*image.size[0]:
                # iterate until the text size is just larger than the criteria
                fontsize += 1
                font = ImageFont.truetype("fonts/Courier New Bold.ttf", fontsize) 
        font = ImageFont.truetype("fonts/Courier New Bold.ttf", fontsize)
        
        # Create white text box
        draw.rectangle((0, 0, width, rect_height), fill='white')

        # Tag image with text
        draw.text((20, rect_height/4),txt, font = font, fill=(0, 0, 0)) 
        
        new_name = fname[:-4] + '_registered.jpg' 
        print(new_name + '\n')
        new_path = tagged_dir + "/" + new_name
        image.save(new_path)
        
    logger.info(f"Tagged Images Saved")

if __name__ == "__main__":
    app()