Az/yolo format support (#619)
* added yolo loader/dumper * changed format_spec * updated reamde, changelog * Used bold font for default dump formatmain
Andrey Zhavoronkov
7 years ago
committed by
Nikita Manovich
parent
f17847ff33
commit
1454ec7ecc
@ -0,0 +1,136 @@
|
||||
format_spec = {
|
||||
"name": "YOLO",
|
||||
"dumpers": [
|
||||
{
|
||||
"display_name": "{name} {format} {version}",
|
||||
"format": "ZIP",
|
||||
"version": "1.0",
|
||||
"handler": "dump"
|
||||
},
|
||||
],
|
||||
"loaders": [
|
||||
{
|
||||
"display_name": "{name} {format} {version}",
|
||||
"format": "ZIP",
|
||||
"version": "1.0",
|
||||
"handler": "load"
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
def get_filename(path):
|
||||
import os
|
||||
return os.path.splitext(os.path.basename(path))[0]
|
||||
|
||||
def load(file_object, annotations):
|
||||
from pyunpack import Archive
|
||||
import os
|
||||
from tempfile import TemporaryDirectory
|
||||
from glob import glob
|
||||
|
||||
def convert_from_yolo(img_size, box):
|
||||
# convertation formulas are based on https://github.com/pjreddie/darknet/blob/master/scripts/voc_label.py
|
||||
# <x> <y> <width> <height> - float values relative to width and height of image
|
||||
# <x> <y> - are center of rectangle
|
||||
def clamp(value, _min, _max):
|
||||
return max(min(_max, value), _min)
|
||||
xtl = clamp(img_size[0] * (box[0] - box[2] / 2), 0, img_size[0])
|
||||
ytl = clamp(img_size[1] * (box[1] - box[3] / 2), 0, img_size[1])
|
||||
xbr = clamp(img_size[0] * (box[0] + box[2] / 2), 0, img_size[0])
|
||||
ybr = clamp(img_size[1] * (box[1] + box[3] / 2), 0, img_size[1])
|
||||
|
||||
return [xtl, ytl, xbr, ybr]
|
||||
|
||||
def parse_yolo_obj(img_size, obj):
|
||||
label_id, x, y, w, h = obj.split(" ")
|
||||
return int(label_id), convert_from_yolo(img_size, (float(x), float(y), float(w), float(h)))
|
||||
|
||||
def match_frame(frame_info, filename):
|
||||
import re
|
||||
# try to match by filename
|
||||
yolo_filename = get_filename(filename)
|
||||
for frame_number, info in frame_info.items():
|
||||
cvat_filename = get_filename(info["path"])
|
||||
if cvat_filename == yolo_filename:
|
||||
return frame_number
|
||||
|
||||
# try to extract frame number from filename
|
||||
numbers = re.findall(r"\d+", filename)
|
||||
if numbers and len(numbers) == 1:
|
||||
return int(numbers[0])
|
||||
|
||||
raise Exception("Cannot match filename or determinate framenumber for {} filename".format(filename))
|
||||
|
||||
def parse_yolo_file(annotation_file, labels_mapping):
|
||||
frame_number = match_frame(annotations.frame_info, annotation_file)
|
||||
with open(annotation_file, "r") as fp:
|
||||
line = fp.readline()
|
||||
while line:
|
||||
frame_info = annotations.frame_info[frame_number]
|
||||
label_id, points = parse_yolo_obj((frame_info["width"], frame_info["height"]), line)
|
||||
annotations.add_shape(annotations.LabeledShape(
|
||||
type="rectangle",
|
||||
frame=frame_number,
|
||||
label=labels_mapping[label_id],
|
||||
points=points,
|
||||
occluded=False,
|
||||
attributes=[],
|
||||
))
|
||||
line = fp.readline()
|
||||
|
||||
def load_labels(labels_file):
|
||||
with open(labels_file, "r") as f:
|
||||
return {idx: label.strip() for idx, label in enumerate(f.readlines()) if label.strip()}
|
||||
|
||||
archive_file = file_object if isinstance(file_object, str) else getattr(file_object, "name")
|
||||
with TemporaryDirectory() as tmp_dir:
|
||||
Archive(archive_file).extractall(tmp_dir)
|
||||
|
||||
labels_file = glob(os.path.join(tmp_dir, "*.names"))
|
||||
if not labels_file:
|
||||
raise Exception("Could not find '*.names' file with labels in uploaded archive")
|
||||
elif len(labels_file) == 1:
|
||||
labels_mapping = load_labels(labels_file[0])
|
||||
else:
|
||||
raise Exception("Too many '*.names' files in uploaded archive: {}".format(labels_file))
|
||||
|
||||
for dirpath, _, filenames in os.walk(tmp_dir):
|
||||
for file in filenames:
|
||||
if ".txt" == os.path.splitext(file)[1]:
|
||||
parse_yolo_file(os.path.join(dirpath, file), labels_mapping)
|
||||
|
||||
def dump(file_object, annotations):
|
||||
from zipfile import ZipFile
|
||||
|
||||
# convertation formulas are based on https://github.com/pjreddie/darknet/blob/master/scripts/voc_label.py
|
||||
# <x> <y> <width> <height> - float values relative to width and height of image
|
||||
# <x> <y> - are center of rectangle
|
||||
def convert_to_yolo(img_size, box):
|
||||
x = (box[0] + box[2]) / 2 / img_size[0]
|
||||
y = (box[1] + box[3]) / 2 / img_size[1]
|
||||
w = (box[2] - box[0]) / img_size[0]
|
||||
h = (box[3] - box[1]) / img_size[1]
|
||||
|
||||
return x, y, w, h
|
||||
|
||||
labels_ids = {label[1]["name"]: idx for idx, label in enumerate(annotations.meta["task"]["labels"])}
|
||||
|
||||
with ZipFile(file_object, "w") as output_zip:
|
||||
for frame_annotation in annotations.group_by_frame():
|
||||
image_name = frame_annotation.name
|
||||
annotation_name = "{}.txt".format(get_filename(image_name))
|
||||
width = frame_annotation.width
|
||||
height = frame_annotation.height
|
||||
|
||||
yolo_annotation = ""
|
||||
for shape in frame_annotation.labeled_shapes:
|
||||
if shape.type != "rectangle":
|
||||
continue
|
||||
|
||||
label = shape.label
|
||||
yolo_bb = convert_to_yolo((width, height), shape.points)
|
||||
yolo_bb = " ".join("{:.6f}".format(p) for p in yolo_bb)
|
||||
yolo_annotation += "{} {}\n".format(labels_ids[label], yolo_bb)
|
||||
|
||||
output_zip.writestr(annotation_name, yolo_annotation)
|
||||
output_zip.writestr("obj.names", "\n".join(l[0] for l in sorted(labels_ids.items(), key=lambda x:x[1])))
|
||||
Loading…
Reference in New Issue