def main(_):
args = default_argument_parser().parse_args()
test_dir = args.test_data_dir
save_dir = args.save_data_dir
wmlu.create_empty_dir(save_dir, remove_if_exists=True)
files = glob.glob(os.path.join(test_dir, "*.jpg"))
m = predm.PredictModel()
m.restoreVariables()
m.remove_batch()
def id_to_text(id):
return m.trainer.category_index[id]
for file in files:
img = wmli.imread(file)
img = np.expand_dims(img, axis=0)
m.predictImages(img)
save_path = os.path.join(save_dir, os.path.basename(file))
xml_path = wmlu.change_suffix(save_path, "xml")
shutil.copy(file, save_path)
labels = [id_to_text(id) for id in m.res_data[RD_LABELS]]
pvt.writeVOCXml(xml_path, m.res_data[RD_BOXES], labels)
if RD_FULL_SIZE_MASKS in m.res_data:
annotations = lmt.trans_odresult_to_annotations_list(m.res_data)
json_path = wmlu.change_suffix(save_path, "json")
lmt.save_labelme_datav1(json_path,
file,
img,
annotations,
label_to_text=id_to_text)
img_save_path = wmlu.change_suffix(xml_path, "jpg")
wmli.imwrite(img_save_path, m.res_data[RD_RESULT_IMAGE])
def get_files(data_dir, img_suffix="jpg"):
files = wmlu.recurse_get_filepath_in_dir(data_dir, suffix=".json")
res = []
for file in files:
img_file = wmlu.change_suffix(file, img_suffix)
img_file1 = wmlu.change_suffix(file, "jpeg")
if os.path.exists(img_file):
res.append((img_file, file))
elif os.path.exists(img_file1):
res.append((img_file1, file))
return res
def trans_one_file(xml_file,save_dir,labels,img_suffix):
img_file = wmlu.change_suffix(xml_file,img_suffix)
if not osp.exists(img_file):
print(f"Find {img_file} faild.")
return
shape, bboxes, labels_names, difficult, truncated,probs = read_voc_xml(xml_file,absolute_coord=True)
_bboxes = []
_labels_name = []
remove_nr = 0
remove_labels = []
bboxes2remove = []
for i,l in enumerate(labels_names):
if l in labels:
remove_nr += 1
bboxes2remove.append(bboxes[i])
remove_labels.append(l)
continue
_bboxes.append(bboxes[i])
_labels_name.append(l)
if remove_nr==0:
wmlu.try_link(img_file, save_dir)
shutil.copy(xml_file,save_dir)
else:
print(f"{wmlu.base_name(xml_file)} remove {remove_nr} labels, labels is {remove_labels}")
img_save_path = osp.join(save_dir,osp.basename(img_file))
xml_save_path = osp.join(save_dir,osp.basename(xml_file))
img = wmli.imread(img_file)
img = wmli.remove_boxes_of_img(img,np.array(bboxes2remove).astype(np.int32))
wmli.imwrite(img_save_path,img)
write_voc_xml(xml_save_path,img_save_path,shape,_bboxes,_labels_name,is_relative_coordinate=False)
def get_shape_from_img(xml_path):
img_path = wmlu.change_suffix(xml_path, "jpg")
if not os.path.exists(img_path):
print(f"Error find img {img_path} faild.")
return [0, 0, 3]
else:
return list(wmli.imread(img_path).shape)
def remove_file(file):
#print(f"Remove {file}.")
#return
os.remove(file)
img_file = wmlu.change_suffix(file,"jpg")
if os.path.exists(img_file):
print(f"Remove {img_file}")
os.remove(img_file)
def getVOCFilesV2(dir_path):
img_files = wmlu.recurse_get_filepath_in_dir(dir_path, ".jpg")
res = []
for f in img_files:
xml_path = wmlu.change_suffix(f, "xml")
if os.path.exists(xml_path) and os.path.exists(f):
res.append([f, xml_path])
return res
def _add_to_tfrecord(self,img_file, tfrecord_writer):
xml_file = wmlu.change_suffix(img_file,"xml")
if not os.path.exists(img_file) or not os.path.exists(xml_file):
print(f"Error file {xml_file}, {img_file}.")
return False
image_data, shape, bboxes, labels, labels_text, difficult, truncated = \
self._process_image(xml_file,img_file)
if image_data is None:
return False
example = self._convert_to_example(image_data, labels, labels_text,
bboxes, shape, difficult, truncated)
tfrecord_writer.write(example.SerializeToString())
return True
def removeUnmatchVOCFiles(dir_path,
image_sub_dir="JPEGImages",
xml_sub_dir="Annotations",
img_suffix=".jpg",
shuffe=False):
if image_sub_dir is not None:
jpeg_dir = os.path.join(dir_path, image_sub_dir)
else:
jpeg_dir = dir_path
if xml_sub_dir is not None:
xml_dir = os.path.join(dir_path, xml_sub_dir)
else:
xml_dir = dir_path
inputfilenames = wml_utils.recurse_get_filepath_in_dir(jpeg_dir,
suffix=img_suffix)
total_removed_jpgs = 0
total_removed_xmls = 0
good_xml_names = []
for file in inputfilenames:
base_name = wmlu.base_name(file)
xml_path = wmlu.change_suffix(file, "xml")
if os.path.exists(xml_path):
good_xml_names.append(base_name)
else:
print(f"remove {file}")
total_removed_jpgs += 1
os.remove(file)
for file in wml_utils.recurse_get_filepath_in_dir(xml_dir, suffix="xml"):
base_name = wmlu.base_name(file)
if base_name not in good_xml_names:
total_removed_xmls += 1
print(f"remove {file}")
os.remove(file)
print(
f"Total remove {total_removed_jpgs} images, total remove {total_removed_xmls} xmls."
)
def get_buddy_file(img_file):
return [wmlu.change_suffix(img_file, "xml")]
def trans_yolotxt(txt_path, classes_names, img_suffix="jpg"):
labels, bboxes = read_yolotxt(txt_path)
img_path = wmlu.change_suffix(txt_path, img_suffix)
xml_path = wmlu.change_suffix(txt_path, "xml")
labels = [classes_names[x] for x in labels]
write_voc_xml(xml_path, img_path, None, bboxes, labels)