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


annotation_classes = AnnotationClass.read_json()
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')