96 lines
3.9 KiB
Java
96 lines
3.9 KiB
Java
package moe.lemonsh.waifu2xlemon;
|
|
|
|
import org.slf4j.Logger;
|
|
import org.slf4j.LoggerFactory;
|
|
|
|
import javax.imageio.ImageIO;
|
|
import java.io.*;
|
|
import java.nio.charset.StandardCharsets;
|
|
import java.nio.file.Files;
|
|
import java.nio.file.Path;
|
|
import java.util.concurrent.BlockingDeque;
|
|
import java.util.concurrent.LinkedBlockingDeque;
|
|
|
|
public class Waifu2x {
|
|
@FunctionalInterface
|
|
public interface WaifuCallback {
|
|
void exec(WaifuResult result);
|
|
}
|
|
|
|
public record WaifuTask(byte[] inputImage, String inputExtension, String outputExtension,
|
|
char noise, char scale, WaifuCallback callback) {}
|
|
|
|
private final BlockingDeque<WaifuTask> taskQueue;
|
|
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();
|
|
taskQueue = new LinkedBlockingDeque<>(2);
|
|
}
|
|
}
|
|
|
|
public synchronized void runQueue() {
|
|
while (true) {
|
|
try {
|
|
WaifuTask task = taskQueue.take();
|
|
log.info("New job! %s image with options noise=%c scale=%c output=%s".formatted(task.inputExtension, task.noise, task.scale, task.outputExtension));
|
|
var result = run(task.inputImage, task.inputExtension, task.outputExtension, task.noise, task.scale);
|
|
log.info("Job completed! Result:\n"+result.stdout);
|
|
task.callback.exec(result);
|
|
} catch (IOException | InterruptedException e) {
|
|
System.exit(10);
|
|
}
|
|
}
|
|
}
|
|
|
|
private boolean imageIllegal(byte[] fullImage) throws IOException {
|
|
// this can most likely be optimized
|
|
var image = ImageIO.read(new ByteArrayInputStream(fullImage));
|
|
return image == null || image.getWidth() > 2000 || image.getHeight() > 2000;
|
|
}
|
|
|
|
public int addTask(WaifuTask task) throws IOException {
|
|
return imageIllegal(task.inputImage) ? 1 : taskQueue.offer(task) ? 0 : 2;
|
|
}
|
|
|
|
static class WaifuResult {
|
|
public boolean success;
|
|
public String stdout;
|
|
public InputStream image;
|
|
public String filename;
|
|
public long size;
|
|
}
|
|
|
|
private WaifuResult run(byte[] 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();
|
|
var result = new WaifuResult();
|
|
result.filename = outputFile.getName();
|
|
var outputBytes = new ByteArrayOutputStream() {
|
|
@Override
|
|
public synchronized byte[] toByteArray() {
|
|
return buf;
|
|
}
|
|
};
|
|
Files.deleteIfExists(inputFile);
|
|
Files.write(inputFile, inputImage);
|
|
var waifuProcess = new ProcessBuilder(ExecutablePath.toString(), "-s", Character.toString(scale), "-n",
|
|
Character.toString(noise), "-i", inputFile.toString(), "-o", outputFilePath.toString()).start();
|
|
waifuProcess.waitFor();
|
|
Files.delete(inputFile);
|
|
result.stdout = new String(waifuProcess.getErrorStream().readAllBytes(), StandardCharsets.UTF_8);
|
|
if (!(result.success = outputFile.isFile())) return result;
|
|
Files.copy(outputFilePath, outputBytes);
|
|
Files.delete(outputFilePath);
|
|
result.image = new ByteArrayInputStream(outputBytes.toByteArray(), 0, outputBytes.size());
|
|
result.size = outputBytes.size();
|
|
return result;
|
|
}
|
|
}
|