91 lines
4.4 KiB
Java
91 lines
4.4 KiB
Java
package moe.lemonsh.waifu2xlemon;
|
|
|
|
import javax.servlet.MultipartConfigElement;
|
|
import javax.servlet.http.Part;
|
|
import java.io.BufferedReader;
|
|
import java.io.IOException;
|
|
import java.io.InputStream;
|
|
import java.io.InputStreamReader;
|
|
import java.net.URLConnection;
|
|
import java.nio.charset.StandardCharsets;
|
|
import java.util.Objects;
|
|
import java.util.stream.Collectors;
|
|
|
|
import static spark.Spark.*;
|
|
|
|
public class App {
|
|
private static final int maxFileSize = 8380416;
|
|
public static void main(String[] argv) throws IOException {
|
|
String envport = System.getenv("WAIFU2XLEMON_PORT");
|
|
String envexec = System.getenv("WAIFU2XLEMON_EXECUTABLE");
|
|
if (envport == null || envexec == null) {
|
|
System.out.println("Error: Invalid config. Make sure that the WAIFU2XLEMON_PORT and WAIFU2XLEMON_EXECUTABLE envvars are set.");
|
|
System.exit(1);
|
|
}
|
|
Waifu2x waifu2x = new Waifu2x(envexec);
|
|
threadPool(4, 1, -1);
|
|
try {
|
|
port(Integer.parseInt(envport));
|
|
} catch (NumberFormatException e) {
|
|
System.out.println("Error: Invalid config. WAIFU2XLEMON_PORT contains an invalid value.");
|
|
System.exit(2);
|
|
}
|
|
MultipartConfigElement multipartConfigElement = new MultipartConfigElement("multipartTemp", maxFileSize, maxFileSize, maxFileSize+1);
|
|
String indexPage;
|
|
try (InputStream indexResource = App.class.getResourceAsStream("/app.html")) {
|
|
indexPage = new BufferedReader(new InputStreamReader(Objects.requireNonNull(indexResource))).lines().collect(Collectors.joining("\n"));
|
|
}
|
|
get("/", (req, resp) -> indexPage);
|
|
post("/", (req, resp) -> {
|
|
req.raw().setAttribute("org.eclipse.jetty.multipartConfig", multipartConfigElement);
|
|
Part filePart, noisePart, typePart, scalePart;
|
|
try {
|
|
filePart = req.raw().getPart("file");
|
|
scalePart = req.raw().getPart("scale");
|
|
noisePart = req.raw().getPart("noise");
|
|
typePart = req.raw().getPart("ftype");
|
|
} catch (IllegalStateException e) {
|
|
return null;
|
|
}
|
|
if (filePart == null || scalePart == null || noisePart == null || typePart == null) halt(400);
|
|
int scale = 0, noise = 0, type = 0;
|
|
try {
|
|
scale = Integer.parseInt(new String(scalePart.getInputStream().readAllBytes(), StandardCharsets.UTF_8));
|
|
noise = Integer.parseInt(new String(noisePart.getInputStream().readAllBytes(), StandardCharsets.UTF_8));
|
|
type = Integer.parseInt(new String(typePart.getInputStream().readAllBytes(), StandardCharsets.UTF_8));
|
|
} catch (NumberFormatException e) {
|
|
halt(400);
|
|
}
|
|
if (scale < 1 || scale > 2 || noise < -1 || noise > 2 || type < 0 || type > 3) halt(400); // boundary check
|
|
String mimeType = URLConnection.guessContentTypeFromName(filePart.getSubmittedFileName());
|
|
String inputExtension, outputExtension = null;
|
|
switch (mimeType) {
|
|
case "image/png" -> inputExtension = ".png";
|
|
case "image/jpeg" -> inputExtension = ".jpg";
|
|
case "image/webp" -> inputExtension = ".webp";
|
|
default -> {
|
|
return "Only PNG/JPG/WEBP files are allowed.";
|
|
}
|
|
}
|
|
switch (type) {
|
|
case 0 -> outputExtension = ".png";
|
|
case 1 -> outputExtension = ".jpg";
|
|
case 2 -> outputExtension = ".webp";
|
|
default -> halt(400);
|
|
}
|
|
Waifu2x.WaifuResult waifuResult;
|
|
waifuResult = waifu2x.run(filePart.getInputStream(), inputExtension, outputExtension, (char)(noise+'0'), (char)(scale+'0'));
|
|
if (waifuResult.success) {
|
|
resp.raw().setContentType(URLConnection.guessContentTypeFromName(waifuResult.filename));
|
|
resp.raw().setHeader("Content-Disposition","attachment; filename="+waifuResult.filename);
|
|
resp.raw().setContentLengthLong(waifuResult.size);
|
|
waifuResult.image.transferTo(resp.raw().getOutputStream());
|
|
resp.raw().getOutputStream().flush();
|
|
resp.raw().getOutputStream().close();
|
|
return null;
|
|
}
|
|
return "Error.\n"+waifuResult.stdout;
|
|
});
|
|
}
|
|
}
|