AI Label

detection_negative_samples.py

@@ -3,16 +3,15 @@ import cv2
 import random
 from PIL import Image
 from lxml import etree
-path_unit= "/data1/zwq/model/data_MPKFNet/data/ConstructionSite/Hatcher_One"
-path_raw = "/data1/zwq/model/data_MPKFNet/data/ConstructionSite/NO"
-result_path = "/data1/zwq/model/data_MPKFNet/data/ConstructionSite/result"
+import argparse,random
+from tqdm import tqdm
 
 class GEN_Annotations:
     def __init__(self, filename):
         self.root = etree.Element("annotation")
 
         child1 = etree.SubElement(self.root, "folder")
-        child1.text = "VOC2007"
+        child1.text = "CDB"
 
         child2 = etree.SubElement(self.root, "filename")
         child2.text = filename
@@ -20,7 +19,7 @@ class GEN_Annotations:
         child3 = etree.SubElement(self.root, "source")
 
         child4 = etree.SubElement(child3, "annotation")
-        child4.text = "PASCAL VOC2007"
+        child4.text = "CDB"
         child5 = etree.SubElement(child3, "database")
         child5.text = "Unknown"
 
@@ -57,16 +56,21 @@ class GEN_Annotations:
         ymaxn.text = str(ymax)
 
 
-def negative_samples():
+def negative_samples(path_unit,path_raw):
     unit_list = os.listdir(path_unit)
+    random.shuffle(unit_list)
+    if len(unit_list) < 1000:
+        unit_list = unit_list
+    else:
+        unit_list = unit_list[:1000]
     unit_img_list = []
     for unit_img in unit_list:
         img1 = Image.open(os.path.join(path_unit,unit_img))
         unit_img_list.append(img1)
     raw_list = os.listdir(path_raw)
-    for raw_path  in raw_list:
-        print(raw_list)
-        print("\n")
+    for raw_path  in tqdm(raw_list):
+        # print(raw_list)
+        # print("\n")
         raw_img = Image.open(os.path.join(path_raw,raw_path))
         unit_num = random.randint(0,10)
         x_left = random.randint(0,raw_img.size[0]-unit_img_list[unit_num].size[0])
@@ -74,15 +78,33 @@ def negative_samples():
         x_right = x_left+unit_img_list[unit_num].size[0]
         y_right = y_left+unit_img_list[unit_num].size[1]
         raw_img.paste(unit_img_list[unit_num],(x_left,y_left,x_right,y_right))
-        raw_img.save(os.path.join(path_raw,raw_path).replace("/NO/","/result/Images/"))
+        raw_img.save(os.path.join(path_raw,raw_path).replace("/NOObject/","/result/Images/"))
         filename = raw_path
         anno = GEN_Annotations(filename)
         anno.set_size(raw_img.size[0], raw_img.size[1], 3)
         anno.add_pic_attr("person", x_left, y_left, x_right, y_right)
-        anno.savefile(os.path.join(path_raw,raw_path).replace("/NO/","/result/Ground_truth/").replace(".jpg",".xml"))
+        anno.savefile(os.path.join(path_raw,raw_path).replace("/NOObject/","/result/Ground_truth/").replace(".jpg",".xml"))
 
 
 
 if __name__ == '__main__':
+    parser1 = argparse.ArgumentParser(description="CUI")
+    parser1.add_argument("path_unit", default='/data/DataManagement/Project/001_cuihu/CAS_CuihuBird/CMB_Cuihubird/FinalData/data',
+                         type=str,
+                         help='path to object unit')
+    parser1.add_argument("path_raw", default='/data/DataManagement/Project/001_cuihu/0009_20220507_RTSP/CDB_CuihuNobird/NOObject', type=str,
+                         help='path to image to deal with')
 
-    negative_samples()
+
+    args = parser1.parse_args()
+    path_unit = args.path_unit
+    path_raw = args.path_raw
+    save_path1 = path_raw.replace("NOObject","result/Images/")
+    print(save_path1)
+    if not os.path.exists(save_path1):
+        os.makedirs(save_path1)
+    save_path2 = path_raw.replace("NOObject","result/Ground_truth/")
+    print(save_path2)
+    if not os.path.exists(save_path2):
+        os.makedirs(save_path2)
+    negative_samples(path_unit,path_raw)