wangchunlin
/
cvat
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Replace YOLO format support in CVAT with Datumaro (#1151)

Browse Source 
* Employ transforms and item wrapper

* Add image class and tests

* Add image info support to formats

* Fix cli

* Fix merge and voc converte

* Update remote images extractor

* Codacy

* Remove item name, require path in Image

* Merge images of dataset items

* Update tests

* Add image dir converter

* Update Datumaro format

* Update COCO format with image info

* Update CVAT format with image info

* Update TFrecord format with image info

* Update VOC formar with image info

* Update YOLO format with image info

* Update dataset manager bindings with image info

* Add image name to id transform

* Replace YOLO export and import in CVAT with Datumaro
main
zhiltsov-max 6 years ago committed by GitHub
parent b36f402f9c
commit 3ce1a6ffe9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 34 additions and 86 deletions
Show Stats Download Patch File Download Diff File

120
cvat/apps/annotation/yolo.py
Unescape Escape View File

@ -24,98 +24,46 @@ format_spec = {
def load(file_object, annotations):
from pyunpack import Archive
import os
import os.path as osp
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 parse_yolo_file(annotation_file, labels_mapping):
frame_number = annotations.match_frame(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()}
from datumaro.plugins.yolo_format.importer import YoloImporter
from cvat.apps.dataset_manager.bindings import import_dm_annotations
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)
image_info = {}
anno_files = glob(osp.join(tmp_dir, '**', '*.txt'), recursive=True)
for filename in anno_files:
filename = osp.basename(filename)
frame_info = None
try:
frame_info = annotations.frame_info[
int(osp.splitext(filename)[0])]
except Exception:
pass
try:
frame_info = annotations.match_frame(filename)
frame_info = annotations.frame_info[frame_info]
except Exception:
pass
if frame_info is not None:
image_info[osp.splitext(filename)[0]] = \
(frame_info['height'], frame_info['width'])
dm_project = YoloImporter()(tmp_dir, image_info=image_info)
dm_dataset = dm_project.make_dataset()
import_dm_annotations(dm_dataset, annotations)
def dump(file_object, annotations):
from zipfile import ZipFile
import os
# 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(os.path.splitext(os.path.basename(image_name))[0])
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])))
from cvat.apps.dataset_manager.bindings import CvatAnnotationsExtractor
from cvat.apps.dataset_manager.util import make_zip_archive
from datumaro.components.project import Environment
from tempfile import TemporaryDirectory
extractor = CvatAnnotationsExtractor('', annotations)
converter = Environment().make_converter('yolo')
with TemporaryDirectory() as temp_dir:
converter(extractor, save_dir=temp_dir)
make_zip_archive(temp_dir, file_object)

Write Preview
Loading…
Cancel
Save