AI Label
This commit is contained in:
1 parent
897af7ce56
commit
ae53e9eac5
6 files changed
+177
-65
No files matched your search
@@ -0,0 +1,30 @@
|
||||
import os,argparse,torch
|
||||
from System_I import Detection_Y5
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser1 = argparse.ArgumentParser(description="AILabel")
|
||||
parser1.add_argument("--RawPath", default='/data/DataManagement/Project/001_cuihu/0012_20220517_RTSP/imagesAILabel', type=str,
|
||||
help='RawPath.')
|
||||
parser1.add_argument("--GPU", default="1", type=str,
|
||||
help='GPU.')
|
||||
parser1.add_argument("--thresh", default=0.4, type=float,
|
||||
help='thresh.')
|
||||
parser1.add_argument("--Weight", default="/data/Hatcher/data/Detection_weigth/BJCHKMCH-Bird-1-4/weights/best.pt", type=str,
|
||||
help='Weight.')
|
||||
args = parser1.parse_args()
|
||||
os.environ['CUDA_VISIBLE_DEVICES'] = args.GPU
|
||||
devices = torch.device('cuda:{}'.format(str(int(args.GPU))) 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","YesResult")):
|
||||
os.makedirs(args.RawPath.replace("imagesAILabel","YesResult"))
|
||||
if not os.path.exists(args.RawPath.replace("imagesAILabel","LabelXML")):
|
||||
os.makedirs(args.RawPath.replace("imagesAILabel","LabelXML"))
|
||||
Detection = Detection_Y5(device=devices,root=" ", detect_thresh=args.thresh,DetWeight=args.Weight)
|
||||
imagepaths = os.listdir(args.RawPath)
|
||||
for imagepath in tqdm(imagepaths):
|
||||
imagepath = os.path.join(args.RawPath,imagepath)
|
||||
Detection.AIlabel(imagepath)
|
||||
Reference in new issue
Block a user