Skip to content

Commit

Permalink
Support for yolo instance-segmentation polygon format (#119)
Browse files Browse the repository at this point in the history 
* added support for yolo instance-segmentation polygon format

* added support for yolo instance-segmentation polygon format

* Update yolo read-write logic, update tests.

---------

Co-authored-by: Markus Fleischhacker <markus.fleischhacker28+git@gmail.com>
  • Loading branch information
@m2key1 @mfl28
m2key1 and mfl28 authored Jul 14, 2024
1 parent 9ff6ea9 commit 3cc364b
Show file tree
Hide file tree
Showing 5 changed files with 512 additions and 486 deletions.
165 changes: 88 additions & 77 deletions src/main/java/com/github/mfl28/boundingboxeditor/model/io/YOLOLoadStrategy.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,8 +38,9 @@
import java.util.stream.Stream;

/**
* Loads rectangular bounding-box annotations in the YOLO-format described at
* <a href="https://github.com/AlexeyAB/Yolo_mark/issues/60#issuecomment-401854885">...</a>
* Loads rectangular bounding-box annotations and instance-segmentation annotations in the YOLO-format described at
* <a href="https://github.com/AlexeyAB/Yolo_mark/issues/60#issuecomment-401854885">...</a> and
* <a href="https://docs.ultralytics.com/datasets/segment/">...</a>
*/
public class YOLOLoadStrategy implements ImageAnnotationLoadStrategy {
public static final String INVALID_BOUNDING_BOX_COORDINATES_MESSAGE = "Invalid bounding-box coordinates on line ";
Expand All @@ -66,18 +67,18 @@ public ImageAnnotationImportResult load(Path path, Set<String> filesToLoad,

try {
loadObjectCategories(path);
} catch(Exception e) {
} catch (Exception e) {
unParsedFileErrorMessages.add(new IOErrorInfoEntry(OBJECT_DATA_FILE_NAME, e.getMessage()));
return new ImageAnnotationImportResult(0, unParsedFileErrorMessages, ImageAnnotationData.empty());
}

if(categories.isEmpty()) {
if (categories.isEmpty()) {
unParsedFileErrorMessages
.add(new IOErrorInfoEntry(OBJECT_DATA_FILE_NAME, "Does not contain any category names."));
return new ImageAnnotationImportResult(0, unParsedFileErrorMessages, ImageAnnotationData.empty());
}

try(Stream<Path> fileStream = Files.walk(path, INCLUDE_SUBDIRECTORIES ? Integer.MAX_VALUE : 1)) {
try (Stream<Path> fileStream = Files.walk(path, INCLUDE_SUBDIRECTORIES ? Integer.MAX_VALUE : 1)) {
List<File> annotationFiles = fileStream
.filter(pathItem -> pathItem.getFileName().toString().endsWith(".txt"))
.map(Path::toFile).toList();
Expand All @@ -86,25 +87,25 @@ public ImageAnnotationImportResult load(Path path, Set<String> filesToLoad,
AtomicInteger nrProcessedFiles = new AtomicInteger(0);

List<ImageAnnotation> imageAnnotations = annotationFiles.parallelStream()
.map(file -> {
progress.set(1.0 * nrProcessedFiles
.incrementAndGet() / totalNrOfFiles);

try {
return loadAnnotationFromFile(file);
} catch(InvalidAnnotationFormatException |
AnnotationToNonExistentImageException |
AnnotationAssociationException |
IOException e) {
unParsedFileErrorMessages
.add(new IOErrorInfoEntry(
file.getName(),
e.getMessage()));
return null;
}
})
.filter(Objects::nonNull)
.toList();
.map(file -> {
progress.set(1.0 * nrProcessedFiles
.incrementAndGet() / totalNrOfFiles);

try {
return loadAnnotationFromFile(file);
} catch (InvalidAnnotationFormatException |
AnnotationToNonExistentImageException |
AnnotationAssociationException |
IOException e) {
unParsedFileErrorMessages
.add(new IOErrorInfoEntry(
file.getName(),
e.getMessage()));
return null;
}
})
.filter(Objects::nonNull)
.toList();

return new ImageAnnotationImportResult(
imageAnnotations.size(),
Expand All @@ -115,18 +116,18 @@ public ImageAnnotationImportResult load(Path path, Set<String> filesToLoad,
}

private void loadObjectCategories(Path root) throws IOException {
if(!root.resolve(OBJECT_DATA_FILE_NAME).toFile().exists()) {
if (!root.resolve(OBJECT_DATA_FILE_NAME).toFile().exists()) {
throw new InvalidAnnotationFormatException(
"Does not exist in annotation folder \"" + root.getFileName().toString() + "\".");
}

try(BufferedReader fileReader = Files.newBufferedReader(root.resolve(OBJECT_DATA_FILE_NAME))) {
try (BufferedReader fileReader = Files.newBufferedReader(root.resolve(OBJECT_DATA_FILE_NAME))) {
String line;

while((line = fileReader.readLine()) != null) {
while ((line = fileReader.readLine()) != null) {
line = line.strip();

if(!line.isBlank()) {
if (!line.isBlank()) {
categories.add(line);
}
}
Expand All @@ -137,38 +138,40 @@ private ImageAnnotation loadAnnotationFromFile(File file) throws IOException {
final List<String> annotatedImageFiles = baseFileNameToImageFileMap.get(
FilenameUtils.getBaseName(file.getName()));

if(annotatedImageFiles == null) {
if (annotatedImageFiles == null) {
throw new AnnotationToNonExistentImageException(
"No associated image file.");
} else if(annotatedImageFiles.size() > 1) {
} else if (annotatedImageFiles.size() > 1) {
throw new AnnotationAssociationException(
"More than one associated image file.");
}

final String annotatedImageFileName = annotatedImageFiles.get(0);
final String annotatedImageFileName = annotatedImageFiles.getFirst();

try(BufferedReader fileReader = Files.newBufferedReader(file.toPath())) {
try (BufferedReader fileReader = Files.newBufferedReader(file.toPath())) {
String line;

List<BoundingShapeData> boundingShapeDataList = new ArrayList<>();

int counter = 1;

while((line = fileReader.readLine()) != null) {
while ((line = fileReader.readLine()) != null) {
line = line.strip();

if(!line.isBlank()) {
if (!line.isBlank()) {
try {
boundingShapeDataList.add(parseBoundingBoxData(line, counter));
} catch(InvalidAnnotationFormatException e) {
final BoundingShapeData boundingShapeData = parseBoundingShapeData(line, counter);
boundingShapeDataList.add(boundingShapeData);
boundingShapeCountPerCategory.merge(boundingShapeData.getCategoryName(), 1, Integer::sum);
} catch (InvalidAnnotationFormatException e) {
unParsedFileErrorMessages.add(new IOErrorInfoEntry(file.getName(), e.getMessage()));
}
}

++counter;
}

if(boundingShapeDataList.isEmpty()) {
if (boundingShapeDataList.isEmpty()) {
return null;
}

Expand All @@ -177,95 +180,103 @@ private ImageAnnotation loadAnnotationFromFile(File file) throws IOException {
}
}

private BoundingBoxData parseBoundingBoxData(String line, int lineNumber) {
private BoundingShapeData parseBoundingShapeData(String line, int lineNumber) {
Scanner scanner = new Scanner(line);
scanner.useLocale(Locale.ENGLISH);

int categoryId = parseCategoryIndex(scanner, lineNumber);

double xMidRelative = parseRatio(scanner, lineNumber);
double yMidRelative = parseRatio(scanner, lineNumber);
double widthRelative = parseRatio(scanner, lineNumber);
double heightRelative = parseRatio(scanner, lineNumber);
List<Double> entries = new ArrayList<>();

while (scanner.hasNextDouble()) {
double entry = scanner.nextDouble();

assertRatio(entry, "Bounds value not within interval [0, 1] on line " + lineNumber + ".");

entries.add(entry);
}

if (entries.size() == 4) {
return createBoundingBoxData(
categoryId, entries.get(0), entries.get(1), entries.get(2), entries.get(3), lineNumber);
} else if(entries.size() >= 6 && entries.size() % 2 == 0) {
return createBoundingPolygonData(categoryId, entries);
}

throw new InvalidAnnotationFormatException("Invalid number of bounds values on line " + lineNumber + ".");
}

private BoundingBoxData createBoundingBoxData(int categoryId, double xMidRelative, double yMidRelative,
double widthRelative, double heightRelative,
int lineNumber) {
double xMinRelative = xMidRelative - widthRelative / 2;
if(xMinRelative < 0 && -xMinRelative < 1e-6) {
if (xMinRelative < 0 && -xMinRelative < 1e-6) {
xMinRelative = 0;
}
assertRatio(xMinRelative, INVALID_BOUNDING_BOX_COORDINATES_MESSAGE + lineNumber + ".");

double yMinRelative = yMidRelative - heightRelative / 2;
if(yMinRelative < 0 && -yMinRelative < 1e-6) {
if (yMinRelative < 0 && -yMinRelative < 1e-6) {
yMinRelative = 0;
}
assertRatio(yMinRelative, INVALID_BOUNDING_BOX_COORDINATES_MESSAGE + lineNumber + ".");

double xMaxRelative = xMidRelative + widthRelative / 2;
if(xMaxRelative > 1 && xMaxRelative - 1 < 1e-6) {
if (xMaxRelative > 1 && xMaxRelative - 1 < 1e-6) {
xMaxRelative = 1;
}
assertRatio(xMaxRelative, INVALID_BOUNDING_BOX_COORDINATES_MESSAGE + lineNumber + ".");

double yMaxRelative = yMidRelative + heightRelative / 2;
if(yMaxRelative > 1 && yMaxRelative - 1 < 1e-6) {
if (yMaxRelative > 1 && yMaxRelative - 1 < 1e-6) {
yMaxRelative = 1;
}
assertRatio(yMaxRelative, INVALID_BOUNDING_BOX_COORDINATES_MESSAGE + lineNumber + ".");

String categoryName = categories.get(categoryId);

ObjectCategory objectCategory = categoryNameToCategoryMap.computeIfAbsent(categoryName,
key -> new ObjectCategory(key,
ColorUtils
.createRandomColor()));
ObjectCategory objectCategory = categoryNameToCategoryMap.computeIfAbsent(
categoryName,
key -> new ObjectCategory(key,
ColorUtils
.createRandomColor()));

// Note that there are no tags or parts in YOLO-format.
BoundingBoxData boundingBoxData = new BoundingBoxData(objectCategory,
xMinRelative, yMinRelative, xMaxRelative, yMaxRelative,
Collections.emptyList());

boundingShapeCountPerCategory.merge(categoryName, 1, Integer::sum);

return boundingBoxData;
return new BoundingBoxData(objectCategory,
xMinRelative, yMinRelative, xMaxRelative, yMaxRelative,
Collections.emptyList());
}

private double parseRatio(Scanner scanner, int lineNumber) {
if(!scanner.hasNextDouble()) {
throw new InvalidAnnotationFormatException(
"Missing or invalid bounding-box bounds on line " + lineNumber + ".");
}

double ratio = scanner.nextDouble();
private BoundingPolygonData createBoundingPolygonData(int categoryId, List<Double> entries) {
String categoryName = categories.get(categoryId);

assertRatio(ratio, lineNumber);
ObjectCategory objectCategory = categoryNameToCategoryMap.computeIfAbsent(categoryName,
key -> new ObjectCategory(key,
ColorUtils
.createRandomColor()));

return ratio;
// Note that there are no tags or parts in YOLO-format.
return new BoundingPolygonData(objectCategory, entries, Collections.emptyList());
}

private int parseCategoryIndex(Scanner scanner, int lineNumber) {
if(!scanner.hasNextInt()) {
if (!scanner.hasNextInt()) {
throw new InvalidAnnotationFormatException("Missing or invalid category index on line " + lineNumber + ".");
}

int categoryId = scanner.nextInt();

if(categoryId < 0 || categoryId >= categories.size()) {
if (categoryId < 0 || categoryId >= categories.size()) {
throw new InvalidAnnotationFormatException("Invalid category index " + categoryId
+ " (of " + categories.size() + " categories) on line " +
lineNumber + ".");
+ " (of " + categories.size() + " categories) on line " +
lineNumber + ".");
}

return categoryId;
}

private void assertRatio(double ratio, int lineNumber) {
if(ratio < 0 || ratio > 1) {
throw new InvalidAnnotationFormatException("Bounds ratio not within [0, 1] on line " + lineNumber + ".");
}
}

private void assertRatio(double ratio, String message) {
if(ratio < 0 || ratio > 1) {
if (ratio < 0 || ratio > 1) {
throw new InvalidAnnotationFormatException(message);
}
}
Expand Down
Loading

0 comments on commit 3cc364b

Please sign in to comment.