import pandas as pd
import numpy as np
import cv2
from src.visualization.description.Annotator import *


def main():
    print("Started...")
    data = pd.read_csv('../../../data/raw/data.csv')
    annotated = pd.read_csv('annotated.csv')
    annotated.dropna(inplace=True)

    print(data.index)
    print(annotated.index)
    data = data.iloc[data.index.difference(annotated.index)]
    print(data.shape)

    data.loc[:, "annotation_0"] = 0
    data.loc[:, "annotation_90"] = 0
    data.loc[:, "annotation_180"] = 0
    data.loc[:, "annotation_270"] = 0

    annotator = Annotator('../../../data/external/pretrained_models/')

    for i, (idx, row) in enumerate(data.iterrows()):
        img_path = f'../../../data/raw/data/{row.uuid}/'
        annotated.loc[idx] = row
        for angle in [0, 90, 180, 270]:
            image = cv2.imread(img_path + f'{angle}.jpg')
            annotation = annotator.annotate(image)
            annotated.loc[idx, f'annotation_{angle}'] = annotation

        if i % 10 == 0:
            print(i, '...')
            annotated.to_csv('annotated.csv', index=False)


    annotated.to_csv('annotated.csv', index=False)

if __name__ == "__main__":
    main()

import pandas as pd
import numpy as np
import cv2
from src.visualization.description.Annotator import *


def main():
    print("Started...")
    data = pd.read_csv('../../../data/raw/data.csv')
    annotated = pd.read_csv('annotated.csv')
    annotated.dropna(inplace=True)

    print(data.index)
    print(annotated.index)
    data = data.iloc[data.index.difference(annotated.index)]
    print(data.shape)

    data.loc[:, "annotation_0"] = 0
    data.loc[:, "annotation_90"] = 0
    data.loc[:, "annotation_180"] = 0
    data.loc[:, "annotation_270"] = 0

    annotator = Annotator('../../../data/external/pretrained_models/')

    for i, (idx, row) in enumerate(data.iterrows()):
        img_path = f'../../../data/raw/data/{row.uuid}/'
        annotated.loc[idx] = row
        for angle in [0, 90, 180, 270]:
            image = cv2.imread(img_path + f'{angle}.jpg')
            annotation = annotator.annotate(image)
            annotated.loc[idx, f'annotation_{angle}'] = annotation

        if i % 10 == 0:
            print(i, '...')
            annotated.to_csv('annotated.csv', index=False)


    annotated.to_csv('annotated.csv', index=False)

if __name__ == "__main__":
    main()
