更新服务器端server.js

4 years ago · becb40d0f4
--- a/style-transform-master/style-transform-master/server.js
+++ b/style-transform-master/style-transform-master/server.js
@ -1,153 +1,211 @@
 'use strict';
 const express = require("express");
 const cookieParser = require("cookie-parser");
 const bodyParser = require('body-parser');
 const cors = require('cors');
 const cheerio = require('cheerio');
 const exec = require('child_process').exec;
 let multer = require('multer');
 const fs = require('fs');
 const crypto = require('crypto')

 var server = express();
 server.use(express.static('public'));
 server.use(cors());
 server.use(bodyParser.json());
 server.use(bodyParser.urlencoded({ extended: false }));
 server.use(cookieParser());

 const port = process.env.PORT || 1337;

 server.get('/', function (req, res) {
 	res.writeHead(200, { 'Content-type': 'text/html' });
 	res.end(fs.readFileSync("./html/index.html"))
 });

 var upload = multer({ dest: 'public/uploads/' });

 var progress = {};
 var relationfile = {};

 var isIdle = true;
 var queue = [];
 var totalTask = 0;

 function run_child(taskID) {
 	let child = exec(`python run_main.py --content ./public/uploads/${relationfile[taskID].content} --style ./public/uploads/${relationfile[taskID].style} --output ./public/uploads/${taskID}.jpg --num_iter ${progress[taskID].max_epoch}`);
 	
 	console.log(taskID + " start!");
 	console.log(`content: ${relationfile[taskID].content}`);
 	console.log(`style: ${relationfile[taskID].style}`)
 	console.log(`output: ${taskID}.jpg`);
 	console.log(`num_iter = ${progress[taskID].max_epoch}`);

 	child.stdout.on("data", function (data) {
 		let args = data.split(' ');
 		let signal = args[0];
 		if (signal != '$') {
 			return;
 		}

 		let iter = parseInt(args[1]);
 		let totalLoss = parseFloat(args[2]);
 		let ContentLoss = parseFloat(args[3]);
 		let StyleLoss = parseFloat(args[4]);
 		let percentage = iter * 100 / progress[taskID].max_epoch;
 		if (percentage > 100) {
 			percentage = 100;
 		}

 		progress[taskID].percentage = percentage;
 		progress[taskID].iter = iter;
 		progress[taskID].totalLoss = totalLoss;
 		progress[taskID].ContentLoss = ContentLoss;
 		progress[taskID].StyleLoss = StyleLoss;
 	});

 	child.on("close", (code) => {
 		if (code == 0) {
 			progress[taskID].status = 1;
 		} else {
 			progress[taskID].status = -1;
 		}
 		console.log(taskID + " done!");
 		totalTask += 1;

 		if (queue.length == 0) {
 			isIdle = true;
 		} else {
 			let nextTaskID = queue.shift();
 			run_child(nextTaskID);
 		}
 	});
 }

 server.post('/generate', upload.fields([{ name: 'content_image', maxCount: 1 }, { name: 'style_image', maxCount: 1 }]), function (req, res) {
 	let max_epoch = parseInt(req.body.max_epoch || "100");
 	let md5 = crypto.createHash('md5');
 	
 	console.log(req.files);
 	console.log("Epoch = " + max_epoch);

 	if (max_epoch > 2000 || max_epoch <= 0 || !req.files.content_image || !req.files.style_image) {
 		res.writeHead(444, { 'Content-type': 'text/html' });
 		res.end(fs.readFileSync("./html/invalidParam.html"));
 		return;
 	}


 	let taskID = Date.now().toString(16) + md5.update(req.files.content_image[0].filename + req.files.style_image[0].filename).digest('hex');

 	relationfile[taskID] = {content: req.files.content_image[0].filename, style: req.files.style_image[0].filename};
 	progress[taskID] = {percentage: 0, max_epoch: max_epoch, status: 0, wait: totalTask + queue.length + 1};

 	if (isIdle) {
 		isIdle = false;
 		progress[taskID].wait -= 1;
 		run_child(taskID);
 	} else {
 		queue.push(taskID);
 	}


 	let $ = cheerio.load(fs.readFileSync("./html/generate.html"));

 	$('script[name="go"]').html(`getProgress("${taskID}"); setInterval(function() {getProgress("${taskID}");}, 1000);`);
 	
 	res.writeHead(200, { 'Content-type': 'text/html' });
 	res.end($.html());
 });

 server.get('/success', function (req, res) {
 	let taskID = req.query.taskID;
 	
 	let $ = cheerio.load(fs.readFileSync("./html/success.html"));

 	$('img[name="content"]').attr("src", `uploads/${relationfile[taskID].content}`);
 	$('img[name="style"]').attr("src", `uploads/${relationfile[taskID].style}`);
 	$('img[name="output"]').attr("src", `uploads/${taskID}.jpg`);

 	$('a[name="content_href"]').attr("href", `uploads/${relationfile[taskID].content}`);
 	$('a[name="style_href"]').attr("href", `uploads/${relationfile[taskID].style}`);
 	$('a[name="output_href"]').attr("href", `uploads/${taskID}.jpg`);

 	res.writeHead(200, { 'Content-type': 'text/html' });
 	res.end($.html());
 });

 server.get('/fail', function (req, res) {
 	res.writeHead(200, { 'Content-type': 'text/html' });
 	res.end(fs.readFileSync("./html/fail.html"));
 });

 server.get('/progress', function (req, res) {
 	if(req.query.taskID && progress[req.query.taskID]) {
 		let result_obj = {};
 		result_obj.percentage = progress[req.query.taskID].percentage;
 		result_obj.wait = progress[req.query.taskID].wait - totalTask;
 		result_obj.status = progress[req.query.taskID].status;
 		res.json(result_obj);
 	}
 });

 server.listen(port);
 'use strict';
 const express = require("express");
 const cookieParser = require("cookie-parser");
 const bodyParser = require('body-parser');
 const cors = require('cors');
 const cheerio = require('cheerio');
 let multer = require('multer');
 const fs = require('fs');
 const crypto = require('crypto');


 const net = require('net');
 const tcp_server = net.createServer();

 var socket = {};
 var SocketConnect = false;

 // 监听 端口
 tcp_server.listen(5678, function () {
 	console.log('tcp_server listening 5678');
 });

 // 处理客户端连接
 tcp_server.on('connection', function (s) {
 	if(SocketConnect) return;
 	console.log(s.address().address + " Connected!");
 	socket = s;
 	DealConnect(s.address().address);
 	SocketConnect = true;
 });

 tcp_server.on('error', function (err) {
 	console.error(err);
 	console.log('tcp_server error!');
 });

 tcp_server.on('close', function () {
 	SocketConnect = false;
 	console.log('tcp_server closed!');
 });

 var callbackGlobal=function(_) {};

 function DealConnect(SocketAddress) {
 	socket.on('data', function (data) {
 		data = data.toString();
 		callbackGlobal(data);
 	});

 	socket.on('close', function () {
 		console.log('client ' + SocketAddress + ' disconneted!');
 	});

 	socket.on("error", function (err) {
 		console.error(err);
 		console.log('client ' + SocketAddress + ' error!');
 	});
 }

 var server = express();
 server.use(express.static('public'));
 server.use(cors());
 server.use(bodyParser.json({limit: "2097152kb"}));
 server.use(bodyParser.urlencoded({ limit: '2097152kb', extended: false }));
 server.use(cookieParser());

 const port = process.env.PORT || 1337;

 server.get('/', function (req, res) {
 	res.writeHead(200, { 'Content-type': 'text/html' });
 	res.end(fs.readFileSync("./html/index.html"))
 });

 var upload = multer({ dest: 'public/uploads/' });

 var progress = {};
 var relationfile = {};

 var isIdle = true;
 var queue = [];
 var totalTask = 0;

 function run_child(taskID) {
 	socket.write(JSON.stringify({
 		taskID: taskID,
 		iter: progress[taskID].max_epoch,
 		content: relationfile[taskID].content,
 		style: relationfile[taskID].style
 	}));
 	callbackGlobal = function (data) {
 		let resultObject = JSON.parse(data);
 		let signal = resultObject.signal;
 		if (signal == '###') {
 			if (resultObject.code == 0) {
 				progress[taskID].status = 1;
 			} else {
 				progress[taskID].status = -1;
 			}
 			console.log(taskID + " done!");
 			totalTask += 1;

 			if (queue.length == 0) {
 				isIdle = true;
 			} else {
 				let nextTaskID = queue.shift();
 				run_child(nextTaskID);
 			}
 		}
 		if (signal != '$') {
 			return;
 		}

 		let iter = parseInt(resultObject.iter);
 		let totalLoss = parseFloat(resultObject.totalLoss);
 		let ContentLoss = parseFloat(resultObject.ContentLoss);
 		let StyleLoss = parseFloat(resultObject.StyleLoss);
 		let percentage = iter * 100 / progress[taskID].max_epoch;
 		if (percentage > 100) {
 			percentage = 100;
 		}

 		progress[taskID].percentage = percentage;
 		progress[taskID].iter = iter;
 		progress[taskID].totalLoss = totalLoss;
 		progress[taskID].ContentLoss = ContentLoss;
 		progress[taskID].StyleLoss = StyleLoss;
 	}
 }

 server.post('/generate', upload.fields([{ name: 'content_image', maxCount: 1 }, { name: 'style_image', maxCount: 1 }]), function (req, res) {
 	let max_epoch = parseInt(req.body.max_epoch || "100");
 	let md5 = crypto.createHash('md5');
 	
 	console.log(req.files);
 	console.log("Epoch = " + max_epoch);

 	if (max_epoch > 2000 || max_epoch <= 0 || !req.files.content_image || !req.files.style_image) {
 		res.writeHead(444, { 'Content-type': 'text/html' });
 		res.end(fs.readFileSync("./html/invalidParam.html"));
 		return;
 	}


 	let taskID = Date.now().toString(16) + md5.update(req.files.content_image[0].filename + req.files.style_image[0].filename).digest('hex');

 	relationfile[taskID] = {content: req.files.content_image[0].filename, style: req.files.style_image[0].filename};
 	progress[taskID] = {percentage: 0, max_epoch: max_epoch, status: 0, wait: totalTask + queue.length + 1};

 	if (isIdle) {
 		isIdle = false;
 		progress[taskID].wait -= 1;
 		run_child(taskID);
 	} else {
 		queue.push(taskID);
 	}


 	let $ = cheerio.load(fs.readFileSync("./html/generate.html"));

 	$('script[name="go"]').html(`getProgress("${taskID}"); setInterval(function() {getProgress("${taskID}");}, 1000);`);
 	
 	res.writeHead(200, { 'Content-type': 'text/html' });
 	res.end($.html());
 });

 server.get('/success', function (req, res) {
 	let taskID = req.query.taskID;
 	if(!relationfile[taskID]) {
 		relationfile[taskID] = {content: "lost.jpg", style: "lost.jpg"};
 	}
 	
 	let $ = cheerio.load(fs.readFileSync("./html/success.html"));

 	$('img[name="content"]').attr("src", `uploads/${relationfile[taskID].content}`);
 	$('img[name="style"]').attr("src", `uploads/${relationfile[taskID].style}`);
 	$('img[name="output"]').attr("src", `uploads/${taskID}.jpg`);

 	$('a[name="content_href"]').attr("href", `uploads/${relationfile[taskID].content}`);
 	$('a[name="style_href"]').attr("href", `uploads/${relationfile[taskID].style}`);
 	$('a[name="output_href"]').attr("href", `uploads/${taskID}.jpg`);

 	res.writeHead(200, { 'Content-type': 'text/html' });
 	res.end($.html());
 });

 server.get('/fail', function (req, res) {
 	res.writeHead(200, { 'Content-type': 'text/html' });
 	res.end(fs.readFileSync("./html/fail.html"));
 });

 server.get('/progress', function (req, res) {
 	if(req.query.taskID && progress[req.query.taskID]) {
 		let result_obj = {};
 		result_obj.percentage = progress[req.query.taskID].percentage;
 		result_obj.wait = progress[req.query.taskID].wait - totalTask;
 		result_obj.status = progress[req.query.taskID].status;
 		res.json(result_obj);
 	}
 });

 server.post('/receive_result', function (req, res) {
 	const taskID = req.body.taskID;
 	const img_base64 = req.body.result;

 	let dataBuffer = new Buffer(img_base64, 'base64');
 	fs.writeFileSync(`./public/uploads/${taskID}.jpg`, dataBuffer);

 	res.json({code: 0});
 });

 server.listen(port);