import os
from pathlib import Path
import cv2
from dto.annotationClass import AnnotationClass
from dto.imageLabel import ImageLabel
from preprocessing import read_labels
from matplotlib import pyplot as plt

from constants import datasets_dir, prefix, processed_images_dir, processed_labels_dir
annotation_classes = AnnotationClass.read_json()


def visualise_dataset():
    cur_dataset = os.path.join(datasets_dir, f'{prefix}2024-06-18', 'train')
    images_dir = os.path.join(cur_dataset, 'images')
    labels_dir = os.path.join(cur_dataset, 'labels')

    for f in os.listdir(images_dir)[35247:]:
        image_path = os.path.join(images_dir, f)
        labels_path = os.path.join(labels_dir, f'{Path(f).stem}.txt')
        img = ImageLabel(
            image_path=image_path,
            image=cv2.imread(image_path),
            labels_path=labels_path,
            labels=read_labels(labels_path)
        )
        img.visualize(annotation_classes)
        print(f'visualizing {image_path}')
        plt.close()
        key = input('Press any key to continue')


def visualise_processed_folder():

    def show_image(img):
        image_path = os.path.join(processed_images_dir, img)
        labels_path = os.path.join(processed_labels_dir, f'{Path(img).stem}.txt')
        img = ImageLabel(
            image_path=image_path,
            image=cv2.imread(image_path),
            labels_path=labels_path,
            labels=read_labels(labels_path)
        )
        img.visualize(annotation_classes)
    images = os.listdir(processed_images_dir)
    cur = 0
    show_image(images[cur])
    pass


if __name__ == '__main__':
    visualise_processed_folder()