66 lines
2.5 KiB
Java
66 lines
2.5 KiB
Java
package moe.lemonsh.waifu2xlemon;
|
|
|
|
import org.slf4j.Logger;
|
|
import org.slf4j.LoggerFactory;
|
|
|
|
import java.io.*;
|
|
import java.nio.charset.StandardCharsets;
|
|
import java.nio.file.Files;
|
|
import java.nio.file.Path;
|
|
|
|
public class Waifu2x {
|
|
public static class WaifuResult {
|
|
public boolean success;
|
|
public String stdout;
|
|
public InputStream image;
|
|
public String filename;
|
|
public long size;
|
|
}
|
|
|
|
private final Object GPULock;
|
|
private final Path ExecutablePath;
|
|
private final Logger log = LoggerFactory.getLogger(Waifu2x.class);
|
|
|
|
public Waifu2x(String wf2xExecutablePath) {
|
|
var executable = new File(wf2xExecutablePath);
|
|
if (!executable.canExecute()) {
|
|
throw new IllegalArgumentException("'%s' is not a valid path.".formatted(wf2xExecutablePath));
|
|
} else {
|
|
ExecutablePath = executable.toPath().toAbsolutePath();
|
|
GPULock = new Object();
|
|
}
|
|
}
|
|
|
|
public WaifuResult run(InputStream inputImage, String inputExtension, String outputExtension, char noise, char scale) throws IOException, InterruptedException {
|
|
var inputFile = new File("wfx_input" + inputExtension).toPath().toAbsolutePath();
|
|
var outputFile = new File("wfx_output" + outputExtension);
|
|
var outputFilePath = outputFile.toPath().toAbsolutePath();
|
|
Files.deleteIfExists(inputFile);
|
|
Files.copy(inputImage, inputFile);
|
|
Process waifuProcess;
|
|
synchronized (GPULock) {
|
|
waifuProcess = new ProcessBuilder(ExecutablePath.toString(), "-s", Character.toString(scale), "-n",
|
|
Character.toString(noise), "-i", inputFile.toString(), "-o", outputFilePath.toString()).start();
|
|
waifuProcess.waitFor();
|
|
}
|
|
var result = new WaifuResult();
|
|
result.filename = outputFile.getName();
|
|
Files.delete(inputFile);
|
|
String stdout = new String(waifuProcess.getErrorStream().readAllBytes(), StandardCharsets.UTF_8);
|
|
log.info(stdout);
|
|
result.stdout = stdout;
|
|
if (!(result.success = outputFile.isFile())) return result;
|
|
var outputBytes = new ByteArrayOutputStream() {
|
|
@Override
|
|
public synchronized byte[] toByteArray() {
|
|
return buf;
|
|
}
|
|
};
|
|
Files.copy(outputFilePath, outputBytes);
|
|
Files.delete(outputFilePath);
|
|
result.image = new ByteArrayInputStream(outputBytes.toByteArray(), 0, outputBytes.size());
|
|
result.size = outputBytes.size();
|
|
return result;
|
|
}
|
|
}
|