40 lines
2.0 KiB
Python
Executable File
40 lines
2.0 KiB
Python
Executable File
import os,argparse,torch
|
|
from System_I import Detection_Y5,Fine_grained_T
|
|
from tqdm import tqdm
|
|
|
|
#os.environ['CUDA_VISIBLE_DEVICES'] = "1"
|
|
|
|
if __name__ == '__main__':
|
|
parser1 = argparse.ArgumentParser(description="AILabel")
|
|
parser1.add_argument("--RawPath", default='/data/DataManagement/Project/006_huanghekou/0006_20220621/Raw_images/imagesAILabel', type=str,
|
|
help='RawPath.')
|
|
parser1.add_argument("--GPU", default="4", type=str,
|
|
help='GPU.')
|
|
parser1.add_argument("--thresh", default=0.4, type=float,
|
|
help='thresh.')
|
|
args = parser1.parse_args()
|
|
os.environ['CUDA_VISIBLE_DEVICES'] = args.GPU#"1"
|
|
try:
|
|
if "imagesAILabel" not in args.RawPath:
|
|
exit()
|
|
else:
|
|
os.system("sudo chmod 777 -R {}".format(str(args.RawPath).replace("imagesAILabel","")))
|
|
except:
|
|
exit()
|
|
#devices = torch.device('cuda:{}'.format(str(int(args.GPU))) if torch.cuda.is_available() else 'cpu')
|
|
devices = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
|
|
if not os.path.exists(args.RawPath.replace("imagesAILabel","NoResult")):
|
|
os.makedirs(args.RawPath.replace("imagesAILabel","NoResult"))
|
|
if not os.path.exists(args.RawPath.replace("imagesAILabel","JPEGImages")):
|
|
os.makedirs(args.RawPath.replace("imagesAILabel","JPEGImages"))
|
|
if not os.path.exists(args.RawPath.replace("imagesAILabel","Annotations")):
|
|
os.makedirs(args.RawPath.replace("imagesAILabel","Annotations"))
|
|
if not os.path.exists(args.RawPath.replace("imagesAILabel", "BrokenImage")):
|
|
os.makedirs(args.RawPath.replace("imagesAILabel", "BrokenImage"))
|
|
|
|
Detection = Detection_Y5(device=devices,root=" ", detect_thresh=args.thresh)
|
|
FineGrained = Fine_grained_T(device=devices)
|
|
imagepaths = os.listdir(args.RawPath)
|
|
for imagepath in tqdm(imagepaths):
|
|
imagepath = os.path.join(args.RawPath,imagepath)
|
|
Detection.AIlabelFine(imagepath,FineGrained) |