110 lines
4.1 KiB
Python
Executable File
110 lines
4.1 KiB
Python
Executable File
import os
|
|
import cv2
|
|
import random
|
|
from PIL import Image
|
|
from lxml import etree
|
|
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 = "CDB"
|
|
|
|
child2 = etree.SubElement(self.root, "filename")
|
|
child2.text = filename
|
|
|
|
child3 = etree.SubElement(self.root, "source")
|
|
|
|
child4 = etree.SubElement(child3, "annotation")
|
|
child4.text = "CDB"
|
|
child5 = etree.SubElement(child3, "database")
|
|
child5.text = "Unknown"
|
|
|
|
child6 = etree.SubElement(child3, "image")
|
|
child6.text = "flickr"
|
|
child7 = etree.SubElement(child3, "flickrid")
|
|
child7.text = "35435"
|
|
|
|
def set_size(self, witdh, height, channel):
|
|
size = etree.SubElement(self.root, "size")
|
|
widthn = etree.SubElement(size, "width")
|
|
widthn.text = str(witdh)
|
|
heightn = etree.SubElement(size, "height")
|
|
heightn.text = str(height)
|
|
channeln = etree.SubElement(size, "depth")
|
|
channeln.text = str(channel)
|
|
|
|
def savefile(self, filename):
|
|
tree = etree.ElementTree(self.root)
|
|
tree.write(filename, pretty_print=True, xml_declaration=False, encoding='utf-8')
|
|
|
|
def add_pic_attr(self, label, xmin, ymin, xmax, ymax):
|
|
object = etree.SubElement(self.root, "object")
|
|
namen = etree.SubElement(object, "name")
|
|
namen.text = label
|
|
bndbox = etree.SubElement(object, "bndbox")
|
|
xminn = etree.SubElement(bndbox, "xmin")
|
|
xminn.text = str(xmin)
|
|
yminn = etree.SubElement(bndbox, "ymin")
|
|
yminn.text = str(ymin)
|
|
xmaxn = etree.SubElement(bndbox, "xmax")
|
|
xmaxn.text = str(xmax)
|
|
ymaxn = etree.SubElement(bndbox, "ymax")
|
|
ymaxn.text = str(ymax)
|
|
|
|
|
|
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 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])
|
|
y_left = random.randint(0,raw_img.size[1]-unit_img_list[unit_num].size[1])
|
|
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("/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("/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')
|
|
|
|
|
|
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) |