Skip to content

Commit 3cc364b

Browse files
m2key1mfl28
m2key1
and
mfl28
authored
Support for yolo instance-segmentation polygon format (#119)
* 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 <[email protected]>
1 parent 9ff6ea9 commit 3cc364b

File tree

5 files changed

+512
-486
lines changed

5 files changed

+512
-486
lines changed

‎src/main/java/com/github/mfl28/boundingboxeditor/model/io/YOLOLoadStrategy.java‎

Lines changed: 88 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -38,8 +38,9 @@
3838
import java.util.stream.Stream;
3939

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

6768
try {
6869
loadObjectCategories(path);
69-
} catch(Exception e) {
70+
} catch (Exception e) {
7071
unParsedFileErrorMessages.add(new IOErrorInfoEntry(OBJECT_DATA_FILE_NAME, e.getMessage()));
7172
return new ImageAnnotationImportResult(0, unParsedFileErrorMessages, ImageAnnotationData.empty());
7273
}
7374

74-
if(categories.isEmpty()) {
75+
if (categories.isEmpty()) {
7576
unParsedFileErrorMessages
7677
.add(new IOErrorInfoEntry(OBJECT_DATA_FILE_NAME, "Does not contain any category names."));
7778
return new ImageAnnotationImportResult(0, unParsedFileErrorMessages, ImageAnnotationData.empty());
7879
}
7980

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

8889
List<ImageAnnotation> imageAnnotations = annotationFiles.parallelStream()
89-
.map(file -> {
90-
progress.set(1.0 * nrProcessedFiles
91-
.incrementAndGet() / totalNrOfFiles);
92-
93-
try {
94-
return loadAnnotationFromFile(file);
95-
} catch(InvalidAnnotationFormatException |
96-
AnnotationToNonExistentImageException |
97-
AnnotationAssociationException |
98-
IOException e) {
99-
unParsedFileErrorMessages
100-
.add(new IOErrorInfoEntry(
101-
file.getName(),
102-
e.getMessage()));
103-
return null;
104-
}
105-
})
106-
.filter(Objects::nonNull)
107-
.toList();
90+
.map(file -> {
91+
progress.set(1.0 * nrProcessedFiles
92+
.incrementAndGet() / totalNrOfFiles);
93+
94+
try {
95+
return loadAnnotationFromFile(file);
96+
} catch (InvalidAnnotationFormatException |
97+
AnnotationToNonExistentImageException |
98+
AnnotationAssociationException |
99+
IOException e) {
100+
unParsedFileErrorMessages
101+
.add(new IOErrorInfoEntry(
102+
file.getName(),
103+
e.getMessage()));
104+
return null;
105+
}
106+
})
107+
.filter(Objects::nonNull)
108+
.toList();
108109

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

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

123-
try(BufferedReader fileReader = Files.newBufferedReader(root.resolve(OBJECT_DATA_FILE_NAME))) {
124+
try (BufferedReader fileReader = Files.newBufferedReader(root.resolve(OBJECT_DATA_FILE_NAME))) {
124125
String line;
125126

126-
while((line = fileReader.readLine()) != null) {
127+
while ((line = fileReader.readLine()) != null) {
127128
line = line.strip();
128129

129-
if(!line.isBlank()) {
130+
if (!line.isBlank()) {
130131
categories.add(line);
131132
}
132133
}
@@ -137,38 +138,40 @@ private ImageAnnotation loadAnnotationFromFile(File file) throws IOException {
137138
final List<String> annotatedImageFiles = baseFileNameToImageFileMap.get(
138139
FilenameUtils.getBaseName(file.getName()));
139140

140-
if(annotatedImageFiles == null) {
141+
if (annotatedImageFiles == null) {
141142
throw new AnnotationToNonExistentImageException(
142143
"No associated image file.");
143-
} else if(annotatedImageFiles.size() > 1) {
144+
} else if (annotatedImageFiles.size() > 1) {
144145
throw new AnnotationAssociationException(
145146
"More than one associated image file.");
146147
}
147148

148-
final String annotatedImageFileName = annotatedImageFiles.get(0);
149+
final String annotatedImageFileName = annotatedImageFiles.getFirst();
149150

150-
try(BufferedReader fileReader = Files.newBufferedReader(file.toPath())) {
151+
try (BufferedReader fileReader = Files.newBufferedReader(file.toPath())) {
151152
String line;
152153

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

155156
int counter = 1;
156157

157-
while((line = fileReader.readLine()) != null) {
158+
while ((line = fileReader.readLine()) != null) {
158159
line = line.strip();
159160

160-
if(!line.isBlank()) {
161+
if (!line.isBlank()) {
161162
try {
162-
boundingShapeDataList.add(parseBoundingBoxData(line, counter));
163-
} catch(InvalidAnnotationFormatException e) {
163+
final BoundingShapeData boundingShapeData = parseBoundingShapeData(line, counter);
164+
boundingShapeDataList.add(boundingShapeData);
165+
boundingShapeCountPerCategory.merge(boundingShapeData.getCategoryName(), 1, Integer::sum);
166+
} catch (InvalidAnnotationFormatException e) {
164167
unParsedFileErrorMessages.add(new IOErrorInfoEntry(file.getName(), e.getMessage()));
165168
}
166169
}
167170

168171
++counter;
169172
}
170173

171-
if(boundingShapeDataList.isEmpty()) {
174+
if (boundingShapeDataList.isEmpty()) {
172175
return null;
173176
}
174177

@@ -177,95 +180,103 @@ private ImageAnnotation loadAnnotationFromFile(File file) throws IOException {
177180
}
178181
}
179182

180-
private BoundingBoxData parseBoundingBoxData(String line, int lineNumber) {
183+
private BoundingShapeData parseBoundingShapeData(String line, int lineNumber) {
181184
Scanner scanner = new Scanner(line);
182185
scanner.useLocale(Locale.ENGLISH);
183186

184187
int categoryId = parseCategoryIndex(scanner, lineNumber);
185188

186-
double xMidRelative = parseRatio(scanner, lineNumber);
187-
double yMidRelative = parseRatio(scanner, lineNumber);
188-
double widthRelative = parseRatio(scanner, lineNumber);
189-
double heightRelative = parseRatio(scanner, lineNumber);
189+
List<Double> entries = new ArrayList<>();
190190

191+
while (scanner.hasNextDouble()) {
192+
double entry = scanner.nextDouble();
193+
194+
assertRatio(entry, "Bounds value not within interval [0, 1] on line " + lineNumber + ".");
195+
196+
entries.add(entry);
197+
}
198+
199+
if (entries.size() == 4) {
200+
return createBoundingBoxData(
201+
categoryId, entries.get(0), entries.get(1), entries.get(2), entries.get(3), lineNumber);
202+
} else if(entries.size() >= 6 && entries.size() % 2 == 0) {
203+
return createBoundingPolygonData(categoryId, entries);
204+
}
205+
206+
throw new InvalidAnnotationFormatException("Invalid number of bounds values on line " + lineNumber + ".");
207+
}
208+
209+
private BoundingBoxData createBoundingBoxData(int categoryId, double xMidRelative, double yMidRelative,
210+
double widthRelative, double heightRelative,
211+
int lineNumber) {
191212
double xMinRelative = xMidRelative - widthRelative / 2;
192-
if(xMinRelative < 0 && -xMinRelative < 1e-6) {
213+
if (xMinRelative < 0 && -xMinRelative < 1e-6) {
193214
xMinRelative = 0;
194215
}
195216
assertRatio(xMinRelative, INVALID_BOUNDING_BOX_COORDINATES_MESSAGE + lineNumber + ".");
196217

197218
double yMinRelative = yMidRelative - heightRelative / 2;
198-
if(yMinRelative < 0 && -yMinRelative < 1e-6) {
219+
if (yMinRelative < 0 && -yMinRelative < 1e-6) {
199220
yMinRelative = 0;
200221
}
201222
assertRatio(yMinRelative, INVALID_BOUNDING_BOX_COORDINATES_MESSAGE + lineNumber + ".");
202223

203224
double xMaxRelative = xMidRelative + widthRelative / 2;
204-
if(xMaxRelative > 1 && xMaxRelative - 1 < 1e-6) {
225+
if (xMaxRelative > 1 && xMaxRelative - 1 < 1e-6) {
205226
xMaxRelative = 1;
206227
}
207228
assertRatio(xMaxRelative, INVALID_BOUNDING_BOX_COORDINATES_MESSAGE + lineNumber + ".");
208229

209230
double yMaxRelative = yMidRelative + heightRelative / 2;
210-
if(yMaxRelative > 1 && yMaxRelative - 1 < 1e-6) {
231+
if (yMaxRelative > 1 && yMaxRelative - 1 < 1e-6) {
211232
yMaxRelative = 1;
212233
}
213234
assertRatio(yMaxRelative, INVALID_BOUNDING_BOX_COORDINATES_MESSAGE + lineNumber + ".");
214235

215236
String categoryName = categories.get(categoryId);
216237

217-
ObjectCategory objectCategory = categoryNameToCategoryMap.computeIfAbsent(categoryName,
218-
key -> new ObjectCategory(key,
219-
ColorUtils
220-
.createRandomColor()));
238+
ObjectCategory objectCategory = categoryNameToCategoryMap.computeIfAbsent(
239+
categoryName,
240+
key -> new ObjectCategory(key,
241+
ColorUtils
242+
.createRandomColor()));
221243

222244
// Note that there are no tags or parts in YOLO-format.
223-
BoundingBoxData boundingBoxData = new BoundingBoxData(objectCategory,
224-
xMinRelative, yMinRelative, xMaxRelative, yMaxRelative,
225-
Collections.emptyList());
226-
227-
boundingShapeCountPerCategory.merge(categoryName, 1, Integer::sum);
228-
229-
return boundingBoxData;
245+
return new BoundingBoxData(objectCategory,
246+
xMinRelative, yMinRelative, xMaxRelative, yMaxRelative,
247+
Collections.emptyList());
230248
}
231249

232-
private double parseRatio(Scanner scanner, int lineNumber) {
233-
if(!scanner.hasNextDouble()) {
234-
throw new InvalidAnnotationFormatException(
235-
"Missing or invalid bounding-box bounds on line " + lineNumber + ".");
236-
}
237-
238-
double ratio = scanner.nextDouble();
250+
private BoundingPolygonData createBoundingPolygonData(int categoryId, List<Double> entries) {
251+
String categoryName = categories.get(categoryId);
239252

240-
assertRatio(ratio, lineNumber);
253+
ObjectCategory objectCategory = categoryNameToCategoryMap.computeIfAbsent(categoryName,
254+
key -> new ObjectCategory(key,
255+
ColorUtils
256+
.createRandomColor()));
241257

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

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

250267
int categoryId = scanner.nextInt();
251268

252-
if(categoryId < 0 || categoryId >= categories.size()) {
269+
if (categoryId < 0 || categoryId >= categories.size()) {
253270
throw new InvalidAnnotationFormatException("Invalid category index " + categoryId
254-
+ " (of " + categories.size() + " categories) on line " +
255-
lineNumber + ".");
271+
+ " (of " + categories.size() + " categories) on line " +
272+
lineNumber + ".");
256273
}
257274

258275
return categoryId;
259276
}
260277

261-
private void assertRatio(double ratio, int lineNumber) {
262-
if(ratio < 0 || ratio > 1) {
263-
throw new InvalidAnnotationFormatException("Bounds ratio not within [0, 1] on line " + lineNumber + ".");
264-
}
265-
}
266-
267278
private void assertRatio(double ratio, String message) {
268-
if(ratio < 0 || ratio > 1) {
279+
if (ratio < 0 || ratio > 1) {
269280
throw new InvalidAnnotationFormatException(message);
270281
}
271282
}

0 commit comments

Comments
 (0)