# Austausch von gelabeldet rgb´s zu ungelabelden rgb´s

import os
import shutil

# Pfade zu den Verzeichnissen
folder_a = "_MOMMERT/sorted_images/good"
new_folder = "_MOMMERT/sorted_images_YOLO_formatted"

# Kopiere folder_a in das aktuelle Verzeichnis und benenne es um
shutil.copytree(folder_a, new_folder)

# Pfade zu den zusätzlichen RGB-Bildverzeichnissen
rgb_folders = [
    "_MOMMERT/wuerzburg_aerial_images_rgb",
    "_MOMMERT/memmingen_aerial_images_rgb",
    "_MOMMERT/konstanz_aerial_images_rgb"
]

# Durchsuche den neuen Ordner nach Bilddateien mit der Endung '_bbox.png'
for filename in os.listdir(new_folder):
    if filename.endswith('_bbox.png'):
        # Ersetze '_bbox.png' durch '.png' im Dateinamen
        new_filename = filename.replace('_bbox.png', '.png')

        # Suche in allen RGB-Verzeichnissen nach der entsprechenden Datei
        file_found = False
        for rgb_folder in rgb_folders:
            file_b_path = os.path.join(rgb_folder, new_filename)

            # Prüfe, ob die Datei in einem der RGB-Ordner vorhanden ist
            if os.path.exists(file_b_path):
                # Pfad zur Originaldatei im neuen Ordner
                original_file_path = os.path.join(new_folder, filename)

                # Kopiere die Datei aus dem RGB-Ordner in den neuen Ordner
                shutil.copy(file_b_path, original_file_path)

                # Neuer Pfadname im neuen Ordner
                new_file_path = os.path.join(new_folder, new_filename)

                # Umbenennen der Datei im neuen Ordner
                os.rename(original_file_path, new_file_path)

                print(f"Kopiert und umbenannt: {file_b_path} -> {new_file_path}")
                
                file_found = True
                break  # Beende die Suche, wenn die Datei gefunden wurde

        if not file_found:
            print(f"Keine passende Datei gefunden für: {new_filename}")

print("Vorgang abgeschlossen.")

# .txt Generator - für die korrekte Konvertierung ins YOLO Format
# Pfad zu Verzeichnis
dir = "_MOMMERT/sorted_images_YOLO_formatted"

# Funktion, um sicherzustellen, dass für jedes Bild eine .txt-Datei vorhanden ist
def create_empty_label_files(directory):
    image_files = [f for f in os.listdir(directory) if f.endswith(('.jpg', '.jpeg', '.png'))]
    for image_file in image_files:
        txt_file = os.path.splitext(image_file)[0] + ".txt"
        txt_path = os.path.join(directory, txt_file)
        if not os.path.exists(txt_path):
            # Erstelle eine leere .txt Datei
            open(txt_path, 'w').close()


create_empty_label_files(dir)

# Sicherstellen, dass alle Bounding-Boxen maximale Werte von 1 und minimale Werte von 0 besitzen
# Pfad zu Ordner 'a', der die .txt-Dateien enthält
folder = "_MOMMERT/sorted_images_YOLO_formatted"

def check_and_correct_labels(label_dir):
    # Durchsuche alle .txt-Dateien im Ordner
    for label_file in os.listdir(label_dir):
        if label_file.endswith('.txt'):
            file_path = os.path.join(label_dir, label_file)
            corrected_lines = []

            with open(file_path, 'r') as f:
                lines = f.readlines()
                for line in lines:
                    # Zerlege die Zeile in einzelne Werte
                    values = line.strip().split()
                    
                    # Erste Zahl ist die Klassen-ID, die bleibt unverändert
                    class_id = values[0]
                    
                    # Konvertiere die restlichen Werte in Fließkommazahlen und prüfe sie
                    coords = [float(v) for v in values[1:]]
                    
                    # Korrigiere die Werte, die außerhalb des Bereichs liegen
                    coords = [min(1.0, max(0.0, v)) for v in coords]
                    
                    # Erstelle die korrigierte Zeile
                    corrected_line = f"{class_id} " + " ".join(map(str, coords)) + "\n"
                    corrected_lines.append(corrected_line)

            # Speichere die korrigierten Koordinaten zurück in die .txt-Datei
            with open(file_path, 'w') as f:
                f.writelines(corrected_lines)

# Beispiel-Aufruf
label_directory = folder  
check_and_correct_labels(label_directory)

