import os
import argparse
import torch
from System_zwq import Detection_Y5, Fine_grained_T
from tqdm import tqdm

#os.environ['CUDA_VISIBLE_DEVICES'] = "1"

def oneFolder():
    RawPath = os.path.join(args.RawPath, "imagesAILabel")
    os.system("sudo chmod 777 -R {}".format(str(args.RawPath)))
    if not os.path.exists(RawPath.replace("imagesAILabel", "NoResult")):
        os.makedirs(RawPath.replace("imagesAILabel", "NoResult"))
    if not os.path.exists(RawPath.replace("imagesAILabel", "JPEGImages")):
        os.makedirs(RawPath.replace("imagesAILabel", "JPEGImages"))
    if not os.path.exists(RawPath.replace("imagesAILabel", "Annotations")):
        os.makedirs(RawPath.replace("imagesAILabel", "Annotations"))
    if not os.path.exists(RawPath.replace("imagesAILabel", "BrokenImage")):
        os.makedirs(RawPath.replace("imagesAILabel", "BrokenImage"))

    imagepaths = os.listdir(RawPath)
    for imagepath in tqdm(imagepaths, ncols=120, ascii=True, desc=RawPath + " images"):
        imagepath = os.path.join(RawPath, imagepath)
        Detection.AIlabelFine(imagepath, FineGrained)


def manyFolder():
    for RawPaths in tqdm(os.listdir(args.RawPath), ncols=120, ascii=True, desc="bird folders"):
    # for RawPaths in os.listdir(args.RawPath):
        RawPath = os.path.join(args.RawPath, RawPaths, "imagesAILabel")
        # os.system("sudo chmod 777 -R {}".format(str(os.path.join(args.RawPath, RawPaths))))
        if not os.path.exists(RawPath.replace("imagesAILabel", "NoResult")):
            os.makedirs(RawPath.replace("imagesAILabel", "NoResult"))
        if not os.path.exists(RawPath.replace("imagesAILabel", "JPEGImages")):
            os.makedirs(RawPath.replace("imagesAILabel", "JPEGImages"))
        if not os.path.exists(RawPath.replace("imagesAILabel", "Annotations")):
            os.makedirs(RawPath.replace("imagesAILabel", "Annotations"))
        if not os.path.exists(RawPath.replace("imagesAILabel", "BrokenImage")):
            os.makedirs(RawPath.replace("imagesAILabel", "BrokenImage"))

        imagepaths = os.listdir(RawPath)
        for imagepath in tqdm(imagepaths, ncols=120, ascii=True, desc=RawPaths + " images"):
        # for imagepath in tqdm(imagepaths):
            imagepath = os.path.join(RawPath, imagepath)
            Detection.AIlabelFine(imagepath, FineGrained)


if __name__ == '__main__':
    parser1 = argparse.ArgumentParser(description="AILabel")
    parser1.add_argument("RawPath", default='/data5/PartDataset/003 PYH/0014 20230411/', type=str, help='RawPath.')
    parser1.add_argument("--GPU", default="7", type=str, help='GPU.')
    parser1.add_argument("--thresh", default=0.4, type=float, help='thresh.')
    args = parser1.parse_args()
    # os.environ['CUDA_VISIBLE_DEVICES'] = "7"  #args.GPU  # "1"

    devices = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    Detection = Detection_Y5(device=devices, root=" ", detect_thresh=args.thresh)
    FineGrained = Fine_grained_T(device=devices)
    oneFolder()
    # manyFolder()