Improve subtitle selection

This commit is contained in:
2023-03-18 18:07:15 +01:00
parent cf64833d3e
commit ba4c1bc1fe
7 changed files with 129 additions and 13 deletions

View File

@@ -42,6 +42,7 @@ public class Config {
private Set<String> forcedKeywords = new HashSet<>(Arrays.asList("forced", "signs", "songs")); private Set<String> forcedKeywords = new HashSet<>(Arrays.asList("forced", "signs", "songs"));
private Set<String> commentaryKeywords = new HashSet<>(Arrays.asList("commentary", "director")); private Set<String> commentaryKeywords = new HashSet<>(Arrays.asList("commentary", "director"));
private Set<String> excludedDirectories = new HashSet<>(); private Set<String> excludedDirectories = new HashSet<>();
private Set<String> preferredSubtitles = new HashSet<>(Arrays.asList("unstyled"));
private List<AttributeConfig> attributeConfig; private List<AttributeConfig> attributeConfig;

View File

@@ -142,7 +142,7 @@ public abstract class ConfigValidator<FieldType> {
protected Predicate<Method> containsGetterOf(ConfigProperty property) { protected Predicate<Method> containsGetterOf(ConfigProperty property) {
return method -> StringUtils.startsWith(method.getName(), "get") return method -> StringUtils.startsWith(method.getName(), "get")
&& StringUtils.containsIgnoreCase(method.getName(), property.prop().replace("-", "")); && StringUtils.equalsIgnoreCase(method.getName().replace("get", ""), property.prop().replace("-", ""));
} }

View File

@@ -15,6 +15,7 @@ import java.io.InputStream;
import java.io.InputStreamReader; import java.io.InputStreamReader;
import java.util.*; import java.util.*;
import java.util.stream.Collectors; import java.util.stream.Collectors;
import java.util.stream.Stream;
import static at.pcgamingfreaks.mkvaudiosubtitlechanger.model.LaneType.AUDIO; import static at.pcgamingfreaks.mkvaudiosubtitlechanger.model.LaneType.AUDIO;
import static at.pcgamingfreaks.mkvaudiosubtitlechanger.model.LaneType.SUBTITLES; import static at.pcgamingfreaks.mkvaudiosubtitlechanger.model.LaneType.SUBTITLES;
@@ -23,6 +24,10 @@ import static java.lang.String.format;
@Slf4j @Slf4j
public class MkvFileProcessor implements FileProcessor { public class MkvFileProcessor implements FileProcessor {
private final ObjectMapper mapper = new ObjectMapper(); private final ObjectMapper mapper = new ObjectMapper();
private static final SubtitleTrackComparator subtitleTrackComparator =
new SubtitleTrackComparator(Config.getInstance().getPreferredSubtitles().toArray(new String[0]));
private static final String DISABLE_DEFAULT_TRACK = "--edit track:%s --set flag-default=0 "; private static final String DISABLE_DEFAULT_TRACK = "--edit track:%s --set flag-default=0 ";
private static final String ENABLE_DEFAULT_TRACK = "--edit track:%s --set flag-default=1 "; private static final String ENABLE_DEFAULT_TRACK = "--edit track:%s --set flag-default=1 ";
private static final String ENABLE_FORCED_TRACK = "--edit track:%s --set flag-forced=1 "; private static final String ENABLE_FORCED_TRACK = "--edit track:%s --set flag-forced=1 ";
@@ -98,23 +103,31 @@ public class MkvFileProcessor implements FileProcessor {
@Override @Override
public void detectDesiredTracks(FileInfoDto info, List<FileAttribute> nonForcedTracks, List<FileAttribute> nonCommentaryTracks, public void detectDesiredTracks(FileInfoDto info, List<FileAttribute> nonForcedTracks, List<FileAttribute> nonCommentaryTracks,
AttributeConfig... configs) { AttributeConfig... configs) {
Set<FileAttribute> tracks = SetUtils.retainOf(nonForcedTracks, nonCommentaryTracks);
Set<FileAttribute> audioTracks = tracks.stream().filter(a -> AUDIO.equals(a.getType())).collect(Collectors.toSet());
Set<FileAttribute> subtitleTracks = tracks.stream().filter(a -> SUBTITLES.equals(a.getType())).collect(Collectors.toSet());
for (AttributeConfig config : configs) { for (AttributeConfig config : configs) {
FileAttribute desiredAudio = null; Optional<FileAttribute> desiredAudio = detectDesiredTrack(config.getAudioLanguage(), audioTracks).findFirst();
FileAttribute desiredSubtitle = null; Optional<FileAttribute> desiredSubtitle = detectDesiredSubtitleTrack(config.getSubtitleLanguage(), subtitleTracks).findFirst();
for (FileAttribute attribute : SetUtils.retainOf(nonForcedTracks, nonCommentaryTracks)) {
if (attribute.getLanguage().equals(config.getAudioLanguage()) if (desiredAudio.isPresent() && desiredSubtitle.isPresent()) {
&& AUDIO.equals(attribute.getType())) desiredAudio = attribute; info.setDesiredAudioLane(desiredAudio.get());
if (attribute.getLanguage().equals(config.getSubtitleLanguage()) info.setDesiredSubtitleLane(desiredSubtitle.get());
&& SUBTITLES.equals(attribute.getType())) desiredSubtitle = attribute;
}
if (desiredAudio != null && desiredSubtitle != null) {
info.setDesiredAudioLane(desiredAudio);
info.setDesiredSubtitleLane(desiredSubtitle);
break; break;
} }
} }
} }
private Stream<FileAttribute> detectDesiredTrack(String language, Set<FileAttribute> tracks) {
return tracks.stream().filter(track -> language.equals(track.getLanguage()));
}
private Stream<FileAttribute> detectDesiredSubtitleTrack(String language, Set<FileAttribute> tracks) {
return detectDesiredTrack(language, tracks)
.sorted(subtitleTrackComparator.reversed());
}
@Override @Override
public List<FileAttribute> retrieveNonForcedTracks(List<FileAttribute> attributes) { public List<FileAttribute> retrieveNonForcedTracks(List<FileAttribute> attributes) {
return attributes.stream() return attributes.stream()

View File

@@ -0,0 +1,34 @@
package at.pcgamingfreaks.mkvaudiosubtitlechanger.impl;
import at.pcgamingfreaks.mkvaudiosubtitlechanger.model.FileAttribute;
import lombok.RequiredArgsConstructor;
import org.apache.commons.lang3.StringUtils;
import java.util.Comparator;
@RequiredArgsConstructor
public class SubtitleTrackComparator implements Comparator<FileAttribute> {
private final String[] preferredSubtitles;
/**
* {@inheritDoc}
*/
@Override
public int compare(FileAttribute track1, FileAttribute track2) {
int result = 0;
if (StringUtils.containsAnyIgnoreCase(track1.getTrackName(), preferredSubtitles)) {
result++;
}
if (StringUtils.containsAnyIgnoreCase(track2.getTrackName(), preferredSubtitles)) {
result--;
}
if (result == 0) {
if (track1.isDefaultTrack()) result++;
if (track2.isDefaultTrack()) result--;
}
return result;
}
}

View File

@@ -4,6 +4,7 @@ import at.pcgamingfreaks.mkvaudiosubtitlechanger.config.Config;
import at.pcgamingfreaks.mkvaudiosubtitlechanger.exceptions.MkvToolNixException; import at.pcgamingfreaks.mkvaudiosubtitlechanger.exceptions.MkvToolNixException;
import at.pcgamingfreaks.mkvaudiosubtitlechanger.impl.FileCollector; import at.pcgamingfreaks.mkvaudiosubtitlechanger.impl.FileCollector;
import at.pcgamingfreaks.mkvaudiosubtitlechanger.impl.FileProcessor; import at.pcgamingfreaks.mkvaudiosubtitlechanger.impl.FileProcessor;
import at.pcgamingfreaks.mkvaudiosubtitlechanger.model.AttributeConfig;
import at.pcgamingfreaks.mkvaudiosubtitlechanger.model.FileAttribute; import at.pcgamingfreaks.mkvaudiosubtitlechanger.model.FileAttribute;
import at.pcgamingfreaks.mkvaudiosubtitlechanger.model.FileInfoDto; import at.pcgamingfreaks.mkvaudiosubtitlechanger.model.FileInfoDto;
import at.pcgamingfreaks.mkvaudiosubtitlechanger.model.ResultStatistic; import at.pcgamingfreaks.mkvaudiosubtitlechanger.model.ResultStatistic;
@@ -93,7 +94,8 @@ public abstract class AttributeUpdaterKernel {
List<FileAttribute> nonCommentaryTracks = processor.retrieveNonCommentaryTracks(attributes); List<FileAttribute> nonCommentaryTracks = processor.retrieveNonCommentaryTracks(attributes);
processor.detectDefaultTracks(fileInfo, attributes, nonForcedTracks); processor.detectDefaultTracks(fileInfo, attributes, nonForcedTracks);
processor.detectDesiredTracks(fileInfo, nonForcedTracks, nonCommentaryTracks); processor.detectDesiredTracks(fileInfo, nonForcedTracks, nonCommentaryTracks,
Config.getInstance().getAttributeConfig().toArray(new AttributeConfig[]{}));
updateFile(fileInfo); updateFile(fileInfo);
} }

View File

@@ -4,6 +4,8 @@ import lombok.AllArgsConstructor;
import lombok.Getter; import lombok.Getter;
import lombok.extern.slf4j.Slf4j; import lombok.extern.slf4j.Slf4j;
import java.util.Objects;
@Slf4j @Slf4j
@Getter @Getter
@AllArgsConstructor @AllArgsConstructor
@@ -15,6 +17,24 @@ public class FileAttribute {
private final boolean forcedTrack; private final boolean forcedTrack;
private final LaneType type; private final LaneType type;
@Override
public boolean equals(Object o) {
if (this == o) return true;
if (o == null || getClass() != o.getClass()) return false;
FileAttribute attribute = (FileAttribute) o;
return id == attribute.id
&& defaultTrack == attribute.defaultTrack
&& forcedTrack == attribute.forcedTrack
&& Objects.equals(language, attribute.language)
&& Objects.equals(trackName, attribute.trackName)
&& type == attribute.type;
}
@Override
public int hashCode() {
return Objects.hash(id, language, trackName, defaultTrack, forcedTrack, type);
}
@Override @Override
public String toString() { public String toString() {
return "[" + "id=" + id + return "[" + "id=" + id +

View File

@@ -0,0 +1,46 @@
package at.pcgamingfreaks.mkvaudiosubtitlechanger.impl;
import at.pcgamingfreaks.mkvaudiosubtitlechanger.model.FileAttribute;
import at.pcgamingfreaks.mkvaudiosubtitlechanger.model.LaneType;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import static org.junit.jupiter.api.Assertions.*;
class SubtitleTrackComparatorTest {
private static final SubtitleTrackComparator comparator = new SubtitleTrackComparator(new String[]{"unstyled"});
private static Stream<Arguments> compareArguments() {
return Stream.of(
Arguments.of(List.of(attr("unstyled sub", false), attr("styled sub", false)),
List.of(attr("unstyled sub", false), attr("styled sub", false))),
Arguments.of(List.of(attr("styled sub", false), attr("unstyled sub", false)),
List.of(attr("unstyled sub", false), attr("styled sub", false))),
Arguments.of(List.of(attr("unstyled sub", true), attr("styled sub", false)),
List.of(attr("unstyled sub", true), attr("styled sub", false))),
Arguments.of(List.of(attr("styled sub", true), attr("unstyled sub", false)),
List.of(attr("unstyled sub", false), attr("styled sub", true))),
Arguments.of(List.of(attr("unstyled sub", true), attr("unstyled sub", false)),
List.of(attr("unstyled sub", true), attr("unstyled sub", false)))
);
}
@ParameterizedTest
@MethodSource("compareArguments")
void compare(List<FileAttribute> input, List<FileAttribute> expected) {
List<FileAttribute> result = input.stream().sorted(comparator.reversed()).collect(Collectors.toList());
assertArrayEquals(expected.toArray(new FileAttribute[0]), result.toArray(new FileAttribute[0]));
}
private static FileAttribute attr(String trackName, boolean defaultTrack) {
return new FileAttribute(0, "", trackName, defaultTrack, false, LaneType.SUBTITLES);
}
}