vouchervision/component_detector/component_detector.py

import os, sys, inspect, json, shutil, cv2, time, glob #imagesize
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
from PIL import Image
from tqdm import tqdm
from time import perf_counter
import concurrent.futures
from threading import Lock
from collections import defaultdict
import multiprocessing
import torch

currentdir = os.path.dirname(inspect.getfile(inspect.currentframe()))
parentdir = os.path.dirname(currentdir)
sys.path.append(currentdir)
from detect import run
sys.path.append(parentdir)
from landmark_processing import LeafSkeleton
from armature_processing import ArmatureSkeleton

def detect_plant_components(cfg, logger, dir_home, Project, Dirs):
    t1_start = perf_counter()
    logger.name = 'Locating Plant Components'
    logger.info(f"Detecting plant components in {len(os.listdir(Project.dir_images))} images")

    try:
        dir_exisiting_labels = cfg['leafmachine']['project']['use_existing_plant_component_detections']
    except:
        dir_exisiting_labels = None
    if cfg['leafmachine']['project']['num_workers'] is None:
        num_workers = 1
    else:
        num_workers = int(cfg['leafmachine']['project']['num_workers'])

    # Weights folder base
    dir_weights = os.path.join(dir_home, 'leafmachine2', 'component_detector','runs','train')
    
    # Detection threshold
    threshold = cfg['leafmachine']['plant_component_detector']['minimum_confidence_threshold']

    detector_version = cfg['leafmachine']['plant_component_detector']['detector_version']
    detector_iteration = cfg['leafmachine']['plant_component_detector']['detector_iteration']
    detector_weights = cfg['leafmachine']['plant_component_detector']['detector_weights']
    weights =  os.path.join(dir_weights,'Plant_Detector',detector_version,detector_iteration,'weights',detector_weights)

    do_save_prediction_overlay_images = not cfg['leafmachine']['plant_component_detector']['do_save_prediction_overlay_images']
    ignore_objects = cfg['leafmachine']['plant_component_detector']['ignore_objects_for_overlay']
    ignore_objects = ignore_objects or []

    if dir_exisiting_labels != None:
        logger.info("Loading existing plant labels")
        fetch_labels(dir_exisiting_labels, os.path.join(Dirs.path_plant_components, 'labels'))
        if len(Project.dir_images) <= 4000:
            logger.debug("Single-threaded create_dictionary_from_txt() len(Project.dir_images) <= 4000")
            A = create_dictionary_from_txt(logger, dir_exisiting_labels, 'Detections_Plant_Components', Project)
        else:
            logger.debug(f"Multi-threaded with ({str(cfg['leafmachine']['project']['num_workers'])}) threads create_dictionary_from_txt() len(Project.dir_images) > 4000")
            A = create_dictionary_from_txt_parallel(logger, cfg, dir_exisiting_labels, 'Detections_Plant_Components', Project)

    else:
        logger.info("Running YOLOv5 to generate plant labels")
        # run(weights = weights,
        #     source = Project.dir_images,
        #     project = Dirs.path_plant_components,
        #     name = Dirs.run_name,
        #     imgsz = (1280, 1280),
        #     nosave = do_save_prediction_overlay_images,
        #     anno_type = 'Plant_Detector',
        #     conf_thres = threshold, 
        #     ignore_objects_for_overlay = ignore_objects,
        #     mode = 'LM2',
        #     LOGGER=logger,)
        source = Project.dir_images
        project = Dirs.path_plant_components
        name = Dirs.run_name
        imgsz = (1280, 1280)
        nosave = do_save_prediction_overlay_images
        anno_type = 'Plant_Detector'
        conf_thres = threshold
        ignore_objects_for_overlay = ignore_objects
        mode = 'LM2'
        LOGGER = logger

        with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:
            futures = [executor.submit(run_in_parallel, weights, source, project, name, imgsz, nosave, anno_type,
                                    conf_thres, 10, ignore_objects_for_overlay, mode, LOGGER, i, num_workers) for i in
                    range(num_workers)]
            for future in concurrent.futures.as_completed(futures):
                try:
                    _ = future.result()
                except Exception as e:
                    logger.error(f'Error in thread: {e}')
                    continue

        t2_stop = perf_counter()
        logger.info(f"[Plant components detection elapsed time] {round(t2_stop - t1_start)} seconds")
        logger.info(f"Threads [{num_workers}]")

        if len(Project.dir_images) <= 4000:
            logger.debug("Single-threaded create_dictionary_from_txt() len(Project.dir_images) <= 4000")
            A = create_dictionary_from_txt(logger, os.path.join(Dirs.path_plant_components, 'labels'), 'Detections_Plant_Components', Project)
        else:
            logger.debug(f"Multi-threaded with ({str(cfg['leafmachine']['project']['num_workers'])}) threads create_dictionary_from_txt() len(Project.dir_images) > 4000")
            A = create_dictionary_from_txt_parallel(logger, cfg, os.path.join(Dirs.path_plant_components, 'labels'), 'Detections_Plant_Components', Project)
    
    dict_to_json(Project.project_data, Dirs.path_plant_components, 'Detections_Plant_Components.json')
    
    t1_stop = perf_counter()
    logger.info(f"[Processing plant components elapsed time] {round(t1_stop - t1_start)} seconds")
    torch.cuda.empty_cache()
    return Project
    

def detect_archival_components(cfg, logger, dir_home, Project, Dirs):
    if not cfg['leafmachine']['use_RGB_label_images']:
        logger.name = 'Skipping LeafMachine2 Label Detection'
        logger.info(f"Full image will be used instead of the label collage")               
    else:
        t1_start = perf_counter()
        logger.name = 'Locating Archival Components'
        logger.info(f"Detecting archival components in {len(os.listdir(Project.dir_images))} images")

        
        try:
            dir_exisiting_labels = cfg['leafmachine']['project']['use_existing_archival_component_detections']
        except:
            dir_exisiting_labels = None
        if cfg['leafmachine']['project']['num_workers'] is None:
            num_workers = 1
        else:
            num_workers = int(cfg['leafmachine']['project']['num_workers'])
        
        # Weights folder base
        dir_weights = os.path.join(dir_home, 'leafmachine2', 'component_detector','runs','train')
        
        # Detection threshold
        threshold = cfg['leafmachine']['archival_component_detector']['minimum_confidence_threshold']

        detector_version = cfg['leafmachine']['archival_component_detector']['detector_version']
        detector_iteration = cfg['leafmachine']['archival_component_detector']['detector_iteration']
        detector_weights = cfg['leafmachine']['archival_component_detector']['detector_weights']
        weights =  os.path.join(dir_weights,'Archival_Detector',detector_version,detector_iteration,'weights',detector_weights)

        do_save_prediction_overlay_images = not cfg['leafmachine']['archival_component_detector']['do_save_prediction_overlay_images']
        ignore_objects = cfg['leafmachine']['archival_component_detector']['ignore_objects_for_overlay']
        ignore_objects = ignore_objects or []


        if dir_exisiting_labels != None:
            logger.info("Loading existing archival labels")
            fetch_labels(dir_exisiting_labels, os.path.join(Dirs.path_archival_components, 'labels'))
            if len(Project.dir_images) <= 4000:
                logger.debug("Single-threaded create_dictionary_from_txt() len(Project.dir_images) <= 4000")
                A = create_dictionary_from_txt(logger, dir_exisiting_labels, 'Detections_Archival_Components', Project)
            else:
                logger.debug(f"Multi-threaded with ({str(cfg['leafmachine']['project']['num_workers'])}) threads create_dictionary_from_txt() len(Project.dir_images) > 4000")
                A = create_dictionary_from_txt_parallel(logger, cfg, dir_exisiting_labels, 'Detections_Archival_Components', Project)

        else:
            logger.info("Running YOLOv5 to generate archival labels")
            # run(weights = weights,
            #     source = Project.dir_images,
            #     project = Dirs.path_archival_components,
            #     name = Dirs.run_name,
            #     imgsz = (1280, 1280),
            #     nosave = do_save_prediction_overlay_images,
            #     anno_type = 'Archival_Detector',
            #     conf_thres = threshold, 
            #     ignore_objects_for_overlay = ignore_objects,
            #     mode = 'LM2',
            #     LOGGER=logger)
            # split the image paths into 4 chunks
            source = Project.dir_images
            project = Dirs.path_archival_components
            name = Dirs.run_name
            imgsz = (1280, 1280)
            nosave = do_save_prediction_overlay_images
            anno_type = 'Archival_Detector'
            conf_thres = threshold
            ignore_objects_for_overlay = ignore_objects
            mode = 'LM2'
            LOGGER = logger

            with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:
                futures = [executor.submit(run_in_parallel, weights, source, project, name, imgsz, nosave, anno_type,
                                        conf_thres, 10, ignore_objects_for_overlay, mode, LOGGER, i, num_workers) for i in
                        range(num_workers)]
                for future in concurrent.futures.as_completed(futures):
                    try:
                        _ = future.result()
                    except Exception as e:
                        logger.error(f'Error in thread: {e}')
                        continue

            t2_stop = perf_counter()
            logger.info(f"[Archival components detection elapsed time] {round(t2_stop - t1_start)} seconds")
            logger.info(f"Threads [{num_workers}]")

            if len(Project.dir_images) <= 4000:
                logger.debug("Single-threaded create_dictionary_from_txt() len(Project.dir_images) <= 4000")
                A = create_dictionary_from_txt(logger, os.path.join(Dirs.path_archival_components, 'labels'), 'Detections_Archival_Components', Project)
            else:
                logger.debug(f"Multi-threaded with ({str(cfg['leafmachine']['project']['num_workers'])}) threads create_dictionary_from_txt() len(Project.dir_images) > 4000")
                A = create_dictionary_from_txt_parallel(logger, cfg, os.path.join(Dirs.path_archival_components, 'labels'), 'Detections_Archival_Components', Project)
        
        dict_to_json(Project.project_data, Dirs.path_archival_components, 'Detections_Archival_Components.json')

        t1_stop = perf_counter()
        logger.info(f"[Processing archival components elapsed time] {round(t1_stop - t1_start)} seconds")
        torch.cuda.empty_cache()
    return Project


def detect_armature_components(cfg, logger, dir_home, Project, Dirs):
    t1_start = perf_counter()
    logger.name = 'Locating Armature Components'
    logger.info(f"Detecting armature components in {len(os.listdir(Project.dir_images))} images")

    if cfg['leafmachine']['project']['num_workers'] is None:
        num_workers = 1
    else:
        num_workers = int(cfg['leafmachine']['project']['num_workers'])

    # Weights folder base
    dir_weights = os.path.join(dir_home, 'leafmachine2', 'component_detector','runs','train')
    
    # Detection threshold
    threshold = cfg['leafmachine']['armature_component_detector']['minimum_confidence_threshold']

    detector_version = cfg['leafmachine']['armature_component_detector']['detector_version']
    detector_iteration = cfg['leafmachine']['armature_component_detector']['detector_iteration']
    detector_weights = cfg['leafmachine']['armature_component_detector']['detector_weights']
    weights =  os.path.join(dir_weights,'Armature_Detector',detector_version,detector_iteration,'weights',detector_weights)

    do_save_prediction_overlay_images = not cfg['leafmachine']['armature_component_detector']['do_save_prediction_overlay_images']
    ignore_objects = cfg['leafmachine']['armature_component_detector']['ignore_objects_for_overlay']
    ignore_objects = ignore_objects or []

    logger.info("Running YOLOv5 to generate armature labels")

    source = Project.dir_images
    project = Dirs.path_armature_components
    name = Dirs.run_name
    imgsz = (1280, 1280)
    nosave = do_save_prediction_overlay_images
    anno_type = 'Armature_Detector'
    conf_thres = threshold
    ignore_objects_for_overlay = ignore_objects
    mode = 'LM2'
    LOGGER = logger

    with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:
        futures = [executor.submit(run_in_parallel, weights, source, project, name, imgsz, nosave, anno_type,
                                conf_thres, 10, ignore_objects_for_overlay, mode, LOGGER, i, num_workers) for i in
                range(num_workers)]
        for future in concurrent.futures.as_completed(futures):
            try:
                _ = future.result()
            except Exception as e:
                logger.error(f'Error in thread: {e}')
                continue

    t2_stop = perf_counter()
    logger.info(f"[Plant components detection elapsed time] {round(t2_stop - t1_start)} seconds")
    logger.info(f"Threads [{num_workers}]")

    if len(Project.dir_images) <= 4000:
        logger.debug("Single-threaded create_dictionary_from_txt() len(Project.dir_images) <= 4000")
        A = create_dictionary_from_txt(logger, os.path.join(Dirs.path_armature_components, 'labels'), 'Detections_Armature_Components', Project)
    else:
        logger.debug(f"Multi-threaded with ({str(cfg['leafmachine']['project']['num_workers'])}) threads create_dictionary_from_txt() len(Project.dir_images) > 4000")
        A = create_dictionary_from_txt_parallel(logger, cfg, os.path.join(Dirs.path_armature_components, 'labels'), 'Detections_Armature_Components', Project)

    dict_to_json(Project.project_data, Dirs.path_armature_components, 'Detections_Armature_Components.json')
    
    t1_stop = perf_counter()
    logger.info(f"[Processing armature components elapsed time] {round(t1_stop - t1_start)} seconds")
    torch.cuda.empty_cache()
    return Project


''' RUN IN PARALLEL'''
def run_in_parallel(weights, source, project, name, imgsz, nosave, anno_type, conf_thres, line_thickness, ignore_objects_for_overlay, mode, LOGGER, chunk, n_workers):
    num_files = len(os.listdir(source))
    LOGGER.info(f"The number of worker threads: ({n_workers}), number of files ({num_files}).")

    chunk_size = len(os.listdir(source)) // n_workers
    start = chunk * chunk_size
    end = start + chunk_size if chunk < (n_workers-1) else len(os.listdir(source))

    sub_source = [os.path.join(source, f) for f in os.listdir(source)[start:end] if f.lower().endswith('.jpg')]

    run(weights=weights,
        source=sub_source,
        project=project,
        name=name,
        imgsz=imgsz,
        nosave=nosave,
        anno_type=anno_type,
        conf_thres=conf_thres,
        ignore_objects_for_overlay=ignore_objects_for_overlay,
        mode=mode,
        LOGGER=LOGGER)

''' RUN IN PARALLEL'''


###### Multi-thread NOTE this works, but unless there are several thousand images, it will be slower
def process_file(logger, file, dir_components, component, Project, lock):
    file_name = str(file.split('.')[0])
    with open(os.path.join(dir_components, file), "r") as f:
        with lock:
            Project.project_data[file_name][component] = [[int(line.split()[0])] + list(map(float, line.split()[1:])) for line in f]
            try:
                image_path = glob.glob(os.path.join(Project.dir_images, file_name + '.*'))[0]
                name_ext = os.path.basename(image_path)
                with Image.open(image_path) as im:
                    _, ext = os.path.splitext(name_ext)
                    if ext not in ['.jpg']:
                        im = im.convert('RGB')
                        im.save(os.path.join(Project.dir_images, file_name) + '.jpg', quality=100)
                        # file_name += '.jpg'
                    width, height = im.size
            except Exception as e:
                print(f"Unable to get image dimensions. Error: {e}")
                logger.info(f"Unable to get image dimensions. Error: {e}")
                width, height = None, None
            if width and height:
                Project.project_data[file_name]['height'] = int(height)
                Project.project_data[file_name]['width'] = int(width)


def create_dictionary_from_txt_parallel(logger, cfg, dir_components, component, Project):
    if cfg['leafmachine']['project']['num_workers'] is None:
        num_workers = 4 
    else:
        num_workers = int(cfg['leafmachine']['project']['num_workers'])

    files = [file for file in os.listdir(dir_components) if file.endswith(".txt")]
    lock = Lock()
    with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:
        futures = []
        for file in files:
            futures.append(executor.submit(process_file, logger, file, dir_components, component, Project, lock))
        for future in concurrent.futures.as_completed(futures):
            pass
    return Project.project_data

######





# Single threaded
def create_dictionary_from_txt(logger, dir_components, component, Project):
    # dict_labels = {}
    for file in tqdm(os.listdir(dir_components), desc="Loading Annotations", colour='green'):
        if file.endswith(".txt"):
            file_name = str(file.split('.')[0])
            with open(os.path.join(dir_components, file), "r") as f:
                # dict_labels[file] = [[int(line.split()[0])] + list(map(float, line.split()[1:])) for line in f]
                Project.project_data[file_name][component] = [[int(line.split()[0])] + list(map(float, line.split()[1:])) for line in f]
                try:
                    image_path = glob.glob(os.path.join(Project.dir_images, file_name + '.*'))[0]
                    name_ext = os.path.basename(image_path)
                    with Image.open(image_path) as im:
                        _, ext = os.path.splitext(name_ext)
                        if ext not in ['.jpg']:
                            im = im.convert('RGB')
                            im.save(os.path.join(Project.dir_images, file_name) + '.jpg', quality=100)
                            # file_name += '.jpg'
                        width, height = im.size
                except Exception as e:
                    # print(f"Unable to get image dimensions. Error: {e}")
                    logger.info(f"Unable to get image dimensions. Error: {e}")
                    width, height = None, None
                if width and height:
                    Project.project_data[file_name]['height'] = int(height)
                    Project.project_data[file_name]['width'] = int(width)
    # for key, value in dict_labels.items():
    #     print(f'{key}  --> {value}')
    return Project.project_data


# old below   
'''def create_dictionary_from_txt(dir_components, component, Project):
    # dict_labels = {}
    for file in os.listdir(dir_components):
        if file.endswith(".txt"):
            file_name = str(file.split('.')[0])
            with open(os.path.join(dir_components, file), "r") as f:
                # dict_labels[file] = [[int(line.split()[0])] + list(map(float, line.split()[1:])) for line in f]
                Project.project_data[file_name][component] = [[int(line.split()[0])] + list(map(float, line.split()[1:])) for line in f]
                try:
                    width, height = imagesize.get(os.path.join(Project.dir_images, '.'.join([file_name,'jpg'])))
                except Exception as e:
                    print(f"Image not in 'jpg' format. Trying 'jpeg'. Note that other formats are not supported.{e}")
                    width, height = imagesize.get(os.path.join(Project.dir_images, '.'.join([file_name,'jpeg'])))
                Project.project_data[file_name]['height'] = int(height)
                Project.project_data[file_name]['width'] = int(width)
    # for key, value in dict_labels.items():
    #     print(f'{key}  --> {value}')
    return Project.project_data'''



def dict_to_json(dict_labels, dir_components, name_json):
    dir_components = os.path.join(dir_components, 'JSON')
    with open(os.path.join(dir_components, name_json), "w") as outfile:
        json.dump(dict_labels, outfile)

def fetch_labels(dir_exisiting_labels, new_dir):
    shutil.copytree(dir_exisiting_labels, new_dir)


'''Landmarks - uses YOLO, but works differently than above. A hybrid between segmentation and component detector'''
def detect_landmarks(cfg, logger, dir_home, Project, batch, n_batches, Dirs, segmentation_complete):
    start_t = perf_counter()
    logger.name = f'[BATCH {batch+1} Detect Landmarks]'
    logger.info(f'Detecting landmarks for batch {batch+1} of {n_batches}')

    landmark_whole_leaves = cfg['leafmachine']['landmark_detector']['landmark_whole_leaves']
    landmark_partial_leaves = cfg['leafmachine']['landmark_detector']['landmark_partial_leaves']

    landmarks_whole_leaves_props = {}
    landmarks_whole_leaves_overlay = {}
    landmarks_partial_leaves_props = {}
    landmarks_partial_leaves_overlay = {}

    if landmark_whole_leaves:
        run_landmarks(cfg, logger, dir_home, Project, batch, n_batches, Dirs, 'Landmarks_Whole_Leaves', segmentation_complete)
    if landmark_partial_leaves:
        run_landmarks(cfg, logger, dir_home, Project, batch, n_batches, Dirs, 'Landmarks_Partial_Leaves', segmentation_complete)

    # if cfg['leafmachine']['leaf_segmentation']['segment_whole_leaves']:
    #     landmarks_whole_leaves_props_batch, landmarks_whole_leaves_overlay_batch = run_landmarks(Instance_Detector_Whole, Project.project_data_list[batch], 0, 
    #                                                                                 "Segmentation_Whole_Leaf", "Whole_Leaf_Cropped", cfg, Project, Dirs, batch, n_batches)#, start+1, end)
    #     landmarks_whole_leaves_props.update(landmarks_whole_leaves_props_batch)
    #     landmarks_whole_leaves_overlay.update(landmarks_whole_leaves_overlay_batch)
    # if cfg['leafmachine']['leaf_segmentation']['segment_partial_leaves']:
    #     landmarks_partial_leaves_props_batch, landmarks_partial_leaves_overlay_batch = run_landmarks(Instance_Detector_Partial, Project.project_data_list[batch], 1, 
    #                                                                                 "Segmentation_Partial_Leaf", "Partial_Leaf_Cropped", cfg, Project, Dirs, batch, n_batches)#, start+1, end)
    #     landmarks_partial_leaves_props.update(landmarks_partial_leaves_props_batch)
    #     landmarks_partial_leaves_overlay.update(landmarks_partial_leaves_overlay_batch)
    
    end_t = perf_counter()
    logger.info(f'Batch {batch+1}/{n_batches}: Landmark Detection Duration --> {round((end_t - start_t)/60)} minutes')
    return Project


def detect_armature(cfg, logger, dir_home, Project, batch, n_batches, Dirs, segmentation_complete):
    start_t = perf_counter()
    logger.name = f'[BATCH {batch+1} Detect Armature]'
    logger.info(f'Detecting armature for batch {batch+1} of {n_batches}')

    landmark_armature = cfg['leafmachine']['modules']['armature']

    landmarks_armature_props = {}
    landmarks_armature_overlay = {}

    if landmark_armature:
        run_armature(cfg, logger, dir_home, Project, batch, n_batches, Dirs, 'Landmarks_Armature', segmentation_complete)

    end_t = perf_counter()
    logger.info(f'Batch {batch+1}/{n_batches}: Armature Detection Duration --> {round((end_t - start_t)/60)} minutes')
    return Project


def run_armature(cfg, logger, dir_home, Project, batch, n_batches, Dirs, leaf_type, segmentation_complete):
    
    logger.info('Detecting armature landmarks from scratch')
    if leaf_type == 'Landmarks_Armature':
        dir_overlay = os.path.join(Dirs.landmarks_armature_overlay, ''.join(['batch_',str(batch+1)]))

    # if not segmentation_complete: # If segmentation was run, then don't redo the unpack, just do the crop into the temp folder
    if leaf_type == 'Landmarks_Armature': # TODO THE 0 is for prickles. For spines I'll need to add a 1 like with partial_leaves or just do it for all
        Project.project_data_list[batch] = unpack_class_from_components_armature(Project.project_data_list[batch], 0, 'Armature_YOLO', 'Armature_BBoxes', Project)
        Project.project_data_list[batch], dir_temp = crop_images_to_bbox_armature(Project.project_data_list[batch], 0, 'Armature_Cropped', "Armature_BBoxes", Project, Dirs, True, cfg)

    # Weights folder base
    dir_weights = os.path.join(dir_home, 'leafmachine2', 'component_detector','runs','train')
    
    # Detection threshold
    threshold = cfg['leafmachine']['landmark_detector_armature']['minimum_confidence_threshold']

    detector_version = cfg['leafmachine']['landmark_detector_armature']['detector_version']
    detector_iteration = cfg['leafmachine']['landmark_detector_armature']['detector_iteration']
    detector_weights = cfg['leafmachine']['landmark_detector_armature']['detector_weights']
    weights =  os.path.join(dir_weights,'Landmark_Detector_YOLO',detector_version,detector_iteration,'weights',detector_weights)

    do_save_prediction_overlay_images = not cfg['leafmachine']['landmark_detector_armature']['do_save_prediction_overlay_images']
    ignore_objects = cfg['leafmachine']['landmark_detector_armature']['ignore_objects_for_overlay']
    ignore_objects = ignore_objects or []
    if cfg['leafmachine']['project']['num_workers'] is None:
        num_workers = 1
    else:
        num_workers = int(cfg['leafmachine']['project']['num_workers'])

    has_images = False
    if len(os.listdir(dir_temp)) > 0:
        has_images = True
        source = dir_temp
        project = dir_overlay
        name = Dirs.run_name
        imgsz = (1280, 1280)
        nosave = do_save_prediction_overlay_images
        anno_type = 'Armature_Detector'
        conf_thres = threshold
        line_thickness = 2
        ignore_objects_for_overlay = ignore_objects
        mode = 'Landmark'
        LOGGER = logger

        # Initialize a Lock object to ensure thread safety
        lock = Lock()

        with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:
            futures = [executor.submit(run_in_parallel, weights, source, project, name, imgsz, nosave, anno_type,
                                    conf_thres, line_thickness, ignore_objects_for_overlay, mode, LOGGER, i, num_workers) for i in
                    range(num_workers)]
            for future in concurrent.futures.as_completed(futures):
                try:
                    _ = future.result()
                except Exception as e:
                    logger.error(f'Error in thread: {e}')
                    continue

        with lock:
            if has_images:
                dimensions_dict = get_cropped_dimensions(dir_temp)
                A = add_to_dictionary_from_txt_armature(cfg, logger, Dirs, leaf_type, os.path.join(dir_overlay, 'labels'), leaf_type, Project, dimensions_dict, dir_temp, batch, n_batches)
            else:
                # TODO add empty placeholder to the image data
                pass
    
    # delete the temp dir
    try:
        shutil.rmtree(dir_temp)
    except:
        try:
            time.sleep(5)
            shutil.rmtree(dir_temp)
        except:
            try:
                time.sleep(5)
                shutil.rmtree(dir_temp)
            except:
                pass

    torch.cuda.empty_cache()

    return Project


def run_landmarks(cfg, logger, dir_home, Project, batch, n_batches, Dirs, leaf_type, segmentation_complete):
    use_existing_landmark_detections = cfg['leafmachine']['landmark_detector']['use_existing_landmark_detections']
    
    if use_existing_landmark_detections is None:
        logger.info('Detecting landmarks from scratch')
        if leaf_type == 'Landmarks_Whole_Leaves':
            dir_overlay = os.path.join(Dirs.landmarks_whole_leaves_overlay, ''.join(['batch_',str(batch+1)]))
        elif leaf_type == 'Landmarks_Partial_Leaves':
            dir_overlay = os.path.join(Dirs.landmarks_partial_leaves_overlay, ''.join(['batch_',str(batch+1)]))

        # if not segmentation_complete: # If segmentation was run, then don't redo the unpack, just do the crop into the temp folder
        if leaf_type == 'Landmarks_Whole_Leaves':
            Project.project_data_list[batch] = unpack_class_from_components(Project.project_data_list[batch], 0, 'Whole_Leaf_BBoxes_YOLO', 'Whole_Leaf_BBoxes', Project)
            Project.project_data_list[batch], dir_temp = crop_images_to_bbox(Project.project_data_list[batch], 0, 'Whole_Leaf_Cropped', "Whole_Leaf_BBoxes", Project, Dirs)

        elif leaf_type == 'Landmarks_Partial_Leaves':
            Project.project_data_list[batch] = unpack_class_from_components(Project.project_data_list[batch], 1, 'Partial_Leaf_BBoxes_YOLO', 'Partial_Leaf_BBoxes', Project)
            Project.project_data_list[batch], dir_temp = crop_images_to_bbox(Project.project_data_list[batch], 1, 'Partial_Leaf_Cropped', "Partial_Leaf_BBoxes", Project, Dirs)
        # else:
        #     if leaf_type == 'Landmarks_Whole_Leaves':
        #         Project.project_data_list[batch], dir_temp = crop_images_to_bbox(Project.project_data_list[batch], 0, 'Whole_Leaf_Cropped', "Whole_Leaf_BBoxes", Project, Dirs)
        #     elif leaf_type == 'Landmarks_Partial_Leaves':
        #         Project.project_data_list[batch], dir_temp = crop_images_to_bbox(Project.project_data_list[batch], 1, 'Partial_Leaf_Cropped', "Partial_Leaf_BBoxes", Project, Dirs)

        # Weights folder base
        dir_weights = os.path.join(dir_home, 'leafmachine2', 'component_detector','runs','train')
        
        # Detection threshold
        threshold = cfg['leafmachine']['landmark_detector']['minimum_confidence_threshold']

        detector_version = cfg['leafmachine']['landmark_detector']['detector_version']
        detector_iteration = cfg['leafmachine']['landmark_detector']['detector_iteration']
        detector_weights = cfg['leafmachine']['landmark_detector']['detector_weights']
        weights =  os.path.join(dir_weights,'Landmark_Detector_YOLO',detector_version,detector_iteration,'weights',detector_weights)

        do_save_prediction_overlay_images = not cfg['leafmachine']['landmark_detector']['do_save_prediction_overlay_images']
        ignore_objects = cfg['leafmachine']['landmark_detector']['ignore_objects_for_overlay']
        ignore_objects = ignore_objects or []
        if cfg['leafmachine']['project']['num_workers'] is None:
            num_workers = 1
        else:
            num_workers = int(cfg['leafmachine']['project']['num_workers'])

        has_images = False
        if len(os.listdir(dir_temp)) > 0:
            has_images = True
            # run(weights = weights,
            #     source = dir_temp,
            #     project = dir_overlay,
            #     name = Dirs.run_name,
            #     imgsz = (1280, 1280),
            #     nosave = do_save_prediction_overlay_images,
            #     anno_type = 'Landmark_Detector_YOLO',
            #     conf_thres = threshold, 
            #     line_thickness = 2,
            #     ignore_objects_for_overlay = ignore_objects,
            #     mode = 'Landmark')
            source = dir_temp
            project = dir_overlay
            name = Dirs.run_name
            imgsz = (1280, 1280)
            nosave = do_save_prediction_overlay_images
            anno_type = 'Landmark_Detector'
            conf_thres = threshold
            line_thickness = 2
            ignore_objects_for_overlay = ignore_objects
            mode = 'Landmark'
            LOGGER = logger

            # Initialize a Lock object to ensure thread safety
            lock = Lock()

            with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:
                futures = [executor.submit(run_in_parallel, weights, source, project, name, imgsz, nosave, anno_type,
                                        conf_thres, line_thickness, ignore_objects_for_overlay, mode, LOGGER, i, num_workers) for i in
                        range(num_workers)]
                for future in concurrent.futures.as_completed(futures):
                    try:
                        _ = future.result()
                    except Exception as e:
                        logger.error(f'Error in thread: {e}')
                        continue

            with lock:
                if has_images:
                    dimensions_dict = get_cropped_dimensions(dir_temp)
                    A = add_to_dictionary_from_txt(cfg, logger, Dirs, leaf_type, os.path.join(dir_overlay, 'labels'), leaf_type, Project, dimensions_dict, dir_temp, batch, n_batches)
                else:
                    # TODO add empty placeholder to the image data
                    pass
    else:
        logger.info('Loading existing landmark annotations')
        dir_temp = os.path.join(use_existing_landmark_detections, f'batch_{str(batch+1)}', 'labels')
        dimensions_dict = get_cropped_dimensions(dir_temp)
        A = add_to_dictionary_from_txt(cfg, logger, Dirs, leaf_type, use_existing_landmark_detections, leaf_type, Project, dimensions_dict, dir_temp, batch, n_batches)

    
    # delete the temp dir
    try:
        shutil.rmtree(dir_temp)
    except:
        try:
            time.sleep(5)
            shutil.rmtree(dir_temp)
        except:
            try:
                time.sleep(5)
                shutil.rmtree(dir_temp)
            except:
                pass

    torch.cuda.empty_cache()

    return Project
    '''def add_to_dictionary_from_txt(cfg, Dirs, leaf_type, dir_components, component, Project, dimensions_dict, dir_temp):
    # dict_labels = {}
    for file in os.listdir(dir_components):
        file_name = str(file.split('.')[0])
        file_name_parent = file_name.split('__')[0]
        Project.project_data[file_name_parent][component] = {}

        if file.endswith(".txt"):
            with open(os.path.join(dir_components, file), "r") as f:
                all_points = [[int(line.split()[0])] + list(map(float, line.split()[1:])) for line in f]
                Project.project_data[file_name_parent][component][file_name] = all_points
                
                height = dimensions_dict[file_name][0]
                width = dimensions_dict[file_name][1]

                Leaf_Skeleton = LeafSkeleton(cfg, Dirs, leaf_type, all_points, height, width, dir_temp, file_name)
                QC_add = Leaf_Skeleton.get_QC()'''


    return Project.project_data

def add_to_dictionary_from_txt_armature(cfg, logger, Dirs, leaf_type, dir_components, component, Project, dimensions_dict, dir_temp, batch, n_batches):
    dpi = cfg['leafmachine']['overlay']['overlay_dpi']
    if leaf_type == 'Landmarks_Armature':
        logger.info(f'Detecting landmarks armature')
        pdf_path = os.path.join(Dirs.landmarks_armature_overlay_QC, ''.join(['landmarks_armature_overlay_QC__',str(batch+1), 'of', str(n_batches), '.pdf']))
        pdf_path_final = os.path.join(Dirs.landmarks_armature_overlay_final, ''.join(['landmarks_armature_overlay_final__',str(batch+1), 'of', str(n_batches), '.pdf']))

    ### FINAL
    # dict_labels = {}
    fig = plt.figure(figsize=(8.27, 11.69), dpi=dpi) # A4 size, 300 dpi
    row, col = 0, 0
    with PdfPages(pdf_path_final) as pdf:
        
        

        for file in os.listdir(dir_components):
            file_name = str(file.split('.')[0])
            file_name_parent = file_name.split('__')[0]

            # Project.project_data_list[batch][file_name_parent][component] = []

            if file_name_parent in Project.project_data_list[batch]:

                

                if file.endswith(".txt"):
                    with open(os.path.join(dir_components, file), "r") as f:
                        all_points = [[int(line.split()[0])] + list(map(float, line.split()[1:])) for line in f]
                        # Project.project_data_list[batch][file_name_parent][component][file_name] = all_points

                        height = dimensions_dict[file_name][0]
                        width = dimensions_dict[file_name][1]

                        Armature_Skeleton = ArmatureSkeleton(cfg, logger, Dirs, leaf_type, all_points, height, width, dir_temp, file_name)
                        Project = add_armature_skeleton_to_project(cfg, logger, Project, batch, file_name_parent, component, Dirs, leaf_type, all_points, height, width, dir_temp, file_name, Armature_Skeleton)
                        final_add = cv2.cvtColor(Armature_Skeleton.get_final(), cv2.COLOR_BGR2RGB)

                        # Add image to the current subplot
                        ax = fig.add_subplot(5, 3, row * 3 + col + 1)
                        ax.imshow(final_add)
                        ax.axis('off')

                        col += 1
                        if col == 3:
                            col = 0
                            row += 1
                        if row == 5:
                            row = 0
                            pdf.savefig(fig)  # Save the current page
                            fig = plt.figure(figsize=(8.27, 11.69), dpi=300) # Create a new page
            else:
                pass

        if row != 0 or col != 0:
            pdf.savefig(fig)  # Save the remaining images on the last page

def add_to_dictionary_from_txt(cfg, logger, Dirs, leaf_type, dir_components, component, Project, dimensions_dict, dir_temp, batch, n_batches):
    dpi = cfg['leafmachine']['overlay']['overlay_dpi']
    if leaf_type == 'Landmarks_Whole_Leaves':
        logger.info(f'Detecting landmarks whole leaves')
        pdf_path = os.path.join(Dirs.landmarks_whole_leaves_overlay_QC, ''.join(['landmarks_whole_leaves_overlay_QC__',str(batch+1), 'of', str(n_batches), '.pdf']))
        pdf_path_final = os.path.join(Dirs.landmarks_whole_leaves_overlay_final, ''.join(['landmarks_whole_leaves_overlay_final__',str(batch+1), 'of', str(n_batches), '.pdf']))
    elif leaf_type == 'Landmarks_Partial_Leaves':
        logger.info(f'Detecting landmarks partial leaves')
        pdf_path = os.path.join(Dirs.landmarks_partial_leaves_overlay_QC, ''.join(['landmarks_partial_leaves_overlay_QC__',str(batch+1), 'of', str(n_batches), '.pdf']))
        pdf_path_final = os.path.join(Dirs.landmarks_partial_leaves_overlay_final, ''.join(['landmarks_partial_leaves_overlay_final__',str(batch+1), 'of', str(n_batches), '.pdf']))
    elif leaf_type == 'Landmarks_Armature':
        logger.info(f'Detecting landmarks armature')
        pdf_path = os.path.join(Dirs.landmarks_armature_overlay_QC, ''.join(['landmarks_armature_overlay_QC__',str(batch+1), 'of', str(n_batches), '.pdf']))
        pdf_path_final = os.path.join(Dirs.landmarks_armature_overlay_final, ''.join(['landmarks_armature_overlay_final__',str(batch+1), 'of', str(n_batches), '.pdf']))

    ### FINAL
    # dict_labels = {}
    fig = plt.figure(figsize=(8.27, 11.69), dpi=dpi) # A4 size, 300 dpi
    row, col = 0, 0
    with PdfPages(pdf_path_final) as pdf:
        
        

        for file in os.listdir(dir_components):
            file_name = str(file.split('.')[0])
            file_name_parent = file_name.split('__')[0]

            # Project.project_data_list[batch][file_name_parent][component] = []

            if file_name_parent in Project.project_data_list[batch]:

                

                if file.endswith(".txt"):
                    with open(os.path.join(dir_components, file), "r") as f:
                        all_points = [[int(line.split()[0])] + list(map(float, line.split()[1:])) for line in f]
                        # Project.project_data_list[batch][file_name_parent][component][file_name] = all_points

                        height = dimensions_dict[file_name][0]
                        width = dimensions_dict[file_name][1]

                        Leaf_Skeleton = LeafSkeleton(cfg, logger, Dirs, leaf_type, all_points, height, width, dir_temp, file_name)
                        Project = add_leaf_skeleton_to_project(cfg, logger, Project, batch, file_name_parent, component, Dirs, leaf_type, all_points, height, width, dir_temp, file_name, Leaf_Skeleton)
                        final_add = cv2.cvtColor(Leaf_Skeleton.get_final(), cv2.COLOR_BGR2RGB)

                        # Add image to the current subplot
                        ax = fig.add_subplot(5, 3, row * 3 + col + 1)
                        ax.imshow(final_add)
                        ax.axis('off')

                        col += 1
                        if col == 3:
                            col = 0
                            row += 1
                        if row == 5:
                            row = 0
                            pdf.savefig(fig)  # Save the current page
                            fig = plt.figure(figsize=(8.27, 11.69), dpi=300) # Create a new page
            else:
                pass

        if row != 0 or col != 0:
            pdf.savefig(fig)  # Save the remaining images on the last page

    ### QC
    '''do_save_QC_pdf = False # TODO refine this
    if do_save_QC_pdf:
        # dict_labels = {}
        fig = plt.figure(figsize=(8.27, 11.69), dpi=dpi) # A4 size, 300 dpi
        row, col = 0, 0
        with PdfPages(pdf_path) as pdf:



            for file in os.listdir(dir_components):
                file_name = str(file.split('.')[0])
                file_name_parent = file_name.split('__')[0]

                if file_name_parent in Project.project_data_list[batch]:

                    if file.endswith(".txt"):
                        with open(os.path.join(dir_components, file), "r") as f:
                            all_points = [[int(line.split()[0])] + list(map(float, line.split()[1:])) for line in f]
                            Project.project_data_list[batch][file_name_parent][component][file_name] = all_points

                            height = dimensions_dict[file_name][0]
                            width = dimensions_dict[file_name][1]

                            Leaf_Skeleton = LeafSkeleton(cfg, logger, Dirs, leaf_type, all_points, height, width, dir_temp, file_name)
                            QC_add = cv2.cvtColor(Leaf_Skeleton.get_QC(), cv2.COLOR_BGR2RGB)

                            # Add image to the current subplot
                            ax = fig.add_subplot(5, 3, row * 3 + col + 1)
                            ax.imshow(QC_add)
                            ax.axis('off')

                            col += 1
                            if col == 3:
                                col = 0
                                row += 1
                            if row == 5:
                                row = 0
                                pdf.savefig(fig)  # Save the current page
                                fig = plt.figure(figsize=(8.27, 11.69), dpi=300) # Create a new page
                else:
                    pass

            if row != 0 or col != 0:
                pdf.savefig(fig)  # Save the remaining images on the last page'''


def add_armature_skeleton_to_project(cfg, logger, Project, batch, file_name_parent, component, Dirs, leaf_type, all_points, height, width, dir_temp, file_name, ARM):
    if ARM.is_complete:
        try:
            Project.project_data_list[batch][file_name_parent][component].append({file_name: [{'armature_status': 'complete'}, {'armature': ARM}]})
        except:
            Project.project_data_list[batch][file_name_parent][component] = []
            Project.project_data_list[batch][file_name_parent][component].append({file_name: [{'armature_status': 'complete'}, {'armature': ARM}]})

    else:
        try:
            Project.project_data_list[batch][file_name_parent][component].append({file_name: [{'armature_status': 'incomplete'}, {'armature': ARM}]})
        except:
            Project.project_data_list[batch][file_name_parent][component] = []
            Project.project_data_list[batch][file_name_parent][component].append({file_name: [{'armature_status': 'incomplete'}, {'armature': ARM}]})


    return Project


def add_leaf_skeleton_to_project(cfg, logger, Project, batch, file_name_parent, component, Dirs, leaf_type, all_points, height, width, dir_temp, file_name, LS):

    if LS.is_complete_leaf:
        try:
            Project.project_data_list[batch][file_name_parent][component].append({file_name: [{'landmark_status': 'complete_leaf'}, {'landmarks': LS}]})
        except:
            Project.project_data_list[batch][file_name_parent][component] = []
            Project.project_data_list[batch][file_name_parent][component].append({file_name: [{'landmark_status': 'complete_leaf'}, {'landmarks': LS}]})
        # Project.project_data_list[batch][file_name_parent][component][file_name].update({'landmark_status': 'complete_leaf'})
        # Project.project_data_list[batch][file_name_parent][component][file_name].update({'landmarks': LS})

    elif LS.is_leaf_no_width:
        try:
            Project.project_data_list[batch][file_name_parent][component].append({file_name: [{'landmark_status': 'leaf_no_width'}, {'landmarks': LS}]})
        except:
            Project.project_data_list[batch][file_name_parent][component] = []
            Project.project_data_list[batch][file_name_parent][component].append({file_name: [{'landmark_status': 'leaf_no_width'}, {'landmarks': LS}]})
        # Project.project_data_list[batch][file_name_parent][component][file_name].update({'landmark_status': 'leaf_no_width'})
        # Project.project_data_list[batch][file_name_parent][component][file_name].update({'landmarks': LS})

    else:
        try:
            Project.project_data_list[batch][file_name_parent][component].append({file_name: [{'landmark_status': 'incomplete'}, {'landmarks': LS}]})
        except:
            Project.project_data_list[batch][file_name_parent][component] = []
            Project.project_data_list[batch][file_name_parent][component].append({file_name: [{'landmark_status': 'incomplete'}, {'landmarks': LS}]})

        # Project.project_data_list[batch][file_name_parent][component][file_name].update({'landmark_status': 'incomplete'})
        # Project.project_data_list[batch][file_name_parent][component][file_name].update({'landmarks': LS})

    return Project


'''
self.determine_lamina_length('final') 

# Lamina tip and base
if self.has_lamina_tip:
    cv2.circle(self.image_final, self.lamina_tip, radius=4, color=(0, 255, 0), thickness=2)
    cv2.circle(self.image_final, self.lamina_tip, radius=2, color=(255, 255, 255), thickness=-1)
if self.has_lamina_base:
    cv2.circle(self.image_final, self.lamina_base, radius=4, color=(255, 0, 0), thickness=2)
    cv2.circle(self.image_final, self.lamina_base, radius=2, color=(255, 255, 255), thickness=-1)

# Apex angle
# if self.apex_center != []:
#     cv2.circle(self.image_final, self.apex_center, radius=3, color=(0, 255, 0), thickness=-1)
if self.apex_left != []:
    cv2.circle(self.image_final, self.apex_left, radius=3, color=(255, 0, 0), thickness=-1)
if self.apex_right != []:
    cv2.circle(self.image_final, self.apex_right, radius=3, color=(0, 0, 255), thickness=-1)

# Base angle
# if self.base_center:
#     cv2.circle(self.image_final, self.base_center, radius=3, color=(0, 255, 0), thickness=-1)
if self.base_left:
    cv2.circle(self.image_final, self.base_left, radius=3, color=(255, 0, 0), thickness=-1)
if self.base_right:
    cv2.circle(self.image_final, self.base_right, radius=3, color=(0, 0, 255), thickness=-1)

# Draw line of fit
for point in self.width_infer:


'''









def get_cropped_dimensions(dir_temp):
    dimensions_dict = {}
    for file_name in os.listdir(dir_temp):
        if file_name.endswith(".jpg"):
            img = cv2.imread(os.path.join(dir_temp, file_name))
            height, width, channels = img.shape
            stem = os.path.splitext(file_name)[0]
            dimensions_dict[stem] = (height, width)
    return dimensions_dict

def unpack_class_from_components_armature(dict_big, cls, dict_name_yolo, dict_name_location, Project):
    # Get the dict that contains plant parts, find the whole leaves
    for filename, value in dict_big.items():
        if "Detections_Armature_Components" in value:
            filtered_components = [val for val in value["Detections_Armature_Components"] if val[0] == cls]
            value[dict_name_yolo] = filtered_components

    for filename, value in dict_big.items():
        if "Detections_Armature_Components" in value:
            filtered_components = [val for val in value["Detections_Armature_Components"] if val[0] == cls]
            height = value['height']
            width = value['width']
            converted_list = [[convert_index_to_class_armature(val[0]), int((val[1] * width) - ((val[3] * width) / 2)), 
                                                                int((val[2] * height) - ((val[4] * height) / 2)), 
                                                                int(val[3] * width) + int((val[1] * width) - ((val[3] * width) / 2)), 
                                                                int(val[4] * height) + int((val[2] * height) - ((val[4] * height) / 2))] for val in filtered_components]
            # Verify that the crops are correct
            # img = Image.open(os.path.join(Project., '.'.join([filename,'jpg'])))
            # for d in converted_list:
            #     img_crop = img.crop((d[1], d[2], d[3], d[4]))
            #     img_crop.show() 
            value[dict_name_location] = converted_list
    # print(dict)
    return dict_big

def unpack_class_from_components(dict_big, cls, dict_name_yolo, dict_name_location, Project):
    # Get the dict that contains plant parts, find the whole leaves
    for filename, value in dict_big.items():
        if "Detections_Plant_Components" in value:
            filtered_components = [val for val in value["Detections_Plant_Components"] if val[0] == cls]
            value[dict_name_yolo] = filtered_components

    for filename, value in dict_big.items():
        if "Detections_Plant_Components" in value:
            filtered_components = [val for val in value["Detections_Plant_Components"] if val[0] == cls]
            height = value['height']
            width = value['width']
            converted_list = [[convert_index_to_class(val[0]), int((val[1] * width) - ((val[3] * width) / 2)), 
                                                                int((val[2] * height) - ((val[4] * height) / 2)), 
                                                                int(val[3] * width) + int((val[1] * width) - ((val[3] * width) / 2)), 
                                                                int(val[4] * height) + int((val[2] * height) - ((val[4] * height) / 2))] for val in filtered_components]
            # Verify that the crops are correct
            # img = Image.open(os.path.join(Project., '.'.join([filename,'jpg'])))
            # for d in converted_list:
            #     img_crop = img.crop((d[1], d[2], d[3], d[4]))
            #     img_crop.show() 
            value[dict_name_location] = converted_list
    # print(dict)
    return dict_big


def crop_images_to_bbox_armature(dict_big, cls, dict_name_cropped, dict_from, Project, Dirs, do_upscale=False, cfg=None):
    dir_temp = os.path.join(Dirs.landmarks, 'TEMP_landmarks')
    os.makedirs(dir_temp, exist_ok=True)
    # For each image, iterate through the whole leaves, segment, report data back to dict_plant_components
    for filename, value in dict_big.items():
        value[dict_name_cropped] = []
        if dict_from in value:
            bboxes_whole_leaves = [val for val in value[dict_from] if val[0] == convert_index_to_class_armature(cls)]
            if len(bboxes_whole_leaves) == 0:
                m = str(''.join(['No objects for class ', convert_index_to_class_armature(0), ' were found']))
                # Print_Verbose(cfg, 3, m).plain()
            else:
                try:
                    img = cv2.imread(os.path.join(Project.dir_images, '.'.join([filename,'jpg'])))
                    # img = cv2.imread(os.path.join(Project, '.'.join([filename,'jpg']))) # Testing
                except:
                    img = cv2.imread(os.path.join(Project.dir_images, '.'.join([filename,'jpeg'])))
                    # img = cv2.imread(os.path.join(Project, '.'.join([filename,'jpeg']))) # Testing
                
                for d in bboxes_whole_leaves:
                    # img_crop = img.crop((d[1], d[2], d[3], d[4])) # PIL
                    img_crop = img[d[2]:d[4], d[1]:d[3]]
                    loc = '-'.join([str(d[1]), str(d[2]), str(d[3]), str(d[4])])
                    # value[dict_name_cropped].append({crop_name: img_crop})
                    if do_upscale:
                        upscale_factor = int(cfg['leafmachine']['landmark_detector_armature']['upscale_factor'])
                        if cls == 0:
                            crop_name = '__'.join([filename,f"PRICKLE-{upscale_factor}x",loc])
                        height, width, _ = img_crop.shape
                        img_crop = cv2.resize(img_crop, ((width * upscale_factor), (height * upscale_factor)), interpolation=cv2.INTER_LANCZOS4)
                    else:
                        if cls == 0:
                            crop_name = '__'.join([filename,'PRICKLE',loc])

                    cv2.imwrite(os.path.join(dir_temp, '.'.join([crop_name,'jpg'])), img_crop)
                    # cv2.imshow('img_crop', img_crop)
                    # cv2.waitKey(0)
                    # img_crop.show() # PIL
    return dict_big, dir_temp


def crop_images_to_bbox(dict_big, cls, dict_name_cropped, dict_from, Project, Dirs):
    dir_temp = os.path.join(Dirs.landmarks, 'TEMP_landmarks')
    os.makedirs(dir_temp, exist_ok=True)
    # For each image, iterate through the whole leaves, segment, report data back to dict_plant_components
    for filename, value in dict_big.items():
        value[dict_name_cropped] = []
        if dict_from in value:
            bboxes_whole_leaves = [val for val in value[dict_from] if val[0] == convert_index_to_class(cls)]
            if len(bboxes_whole_leaves) == 0:
                m = str(''.join(['No objects for class ', convert_index_to_class(0), ' were found']))
                # Print_Verbose(cfg, 3, m).plain()
            else:
                try:
                    img = cv2.imread(os.path.join(Project.dir_images, '.'.join([filename,'jpg'])))
                    # img = cv2.imread(os.path.join(Project, '.'.join([filename,'jpg']))) # Testing
                except:
                    img = cv2.imread(os.path.join(Project.dir_images, '.'.join([filename,'jpeg'])))
                    # img = cv2.imread(os.path.join(Project, '.'.join([filename,'jpeg']))) # Testing
                
                for d in bboxes_whole_leaves:
                    # img_crop = img.crop((d[1], d[2], d[3], d[4])) # PIL
                    img_crop = img[d[2]:d[4], d[1]:d[3]]
                    loc = '-'.join([str(d[1]), str(d[2]), str(d[3]), str(d[4])])
                    if cls == 0:
                        crop_name = '__'.join([filename,'L',loc])
                    elif cls == 1:
                        crop_name = '__'.join([filename,'PL',loc])
                    elif cls == 2:
                        crop_name = '__'.join([filename,'ARM',loc])
                    # value[dict_name_cropped].append({crop_name: img_crop})
                    cv2.imwrite(os.path.join(dir_temp, '.'.join([crop_name,'jpg'])), img_crop)
                    # cv2.imshow('img_crop', img_crop)
                    # cv2.waitKey(0)
                    # img_crop.show() # PIL
    return dict_big, dir_temp

def convert_index_to_class(ind):
    mapping = {
        0: 'apex_angle',
        1: 'base_angle',
        2: 'lamina_base',
        3: 'lamina_tip',
        4: 'lamina_width',
        5: 'lobe_tip',
        6: 'midvein_trace',
        7: 'petiole_tip',
        8: 'petiole_trace',
    }
    return mapping.get(ind, 'Invalid class').lower()

def convert_index_to_class_armature(ind):
    mapping = {
        0: 'tip',
        1: 'middle',
        2: 'outer',
    }
    return mapping.get(ind, 'Invalid class').lower()