61 lines
3.0 KiB
Python
Executable File
61 lines
3.0 KiB
Python
Executable File
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()
|
|
|