Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support for yolo instance-segmentation polygon format #119

Merged
merged 4 commits into from
Jul 14, 2024
Merged

Support for yolo instance-segmentation polygon format #119

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
35f37a7
added support for yolo instance-segmentation polygon format
m2key1 Mar 4, 2024
dd87cb8
added support for yolo instance-segmentation polygon format
m2key1 Mar 4, 2024
9774eb5
Merge branch 'master' into yolo-segmentation
mfl28 Jul 13, 2024
298ee4a
Update yolo read-write logic, update tests.
mfl28 Jul 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
165 changes: 88 additions & 77 deletions src/main/java/com/github/mfl28/boundingboxeditor/model/io/YOLOLoadStrategy.java
View file
Open in desktop
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