Skip to content

Commit 79a4b0c

Browse files
589hero589hero
committedFeb 15, 2022
initial commit
0 parents  commit 79a4b0c

14 files changed

+253
-0
lines changed
 

‎Dockerfile

+8
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
FROM 589hero/u2net-onnx:latest
2+
3+
COPY . /app
4+
WORKDIR /app
5+
6+
RUN pip install -r requirements.txt
7+
8+
CMD ["python", "server.py"]

‎requirements.txt

+5
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
Flask==1.1.2
2+
pillow==9.0.0
3+
onnxruntime==1.10.0
4+
rembg==2.0.5
5+
waitress==2.0.0

‎server.py

+93
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
import os
2+
import io
3+
import time
4+
import threading
5+
import onnxruntime as ort
6+
from PIL import Image, ImageFile
7+
from queue import Queue, Empty
8+
from flask import Flask, request, render_template, send_file, jsonify
9+
from rembg.bg import remove
10+
11+
app = Flask(__name__, template_folder='./templates/')
12+
13+
ImageFile.LOAD_TRUNCATED_IMAGES = True
14+
requests_queue = Queue()
15+
BATCH_SIZE = 1
16+
CHECK_INTERVAL = 0.1
17+
u2net_path = os.environ['U2NET_HOME'] + '/u2net.onnx'
18+
ort_session = ort.InferenceSession(u2net_path, providers=ort.get_available_providers())
19+
20+
21+
def handle_requests_by_batch():
22+
while True:
23+
requests_batch = []
24+
while not (len(requests_batch) >= BATCH_SIZE):
25+
try:
26+
requests_batch.append(requests_queue.get(timeout=CHECK_INTERVAL))
27+
except Empty:
28+
continue
29+
30+
for request in requests_batch:
31+
org_img_bytes, bg_img_bytes = request['input']
32+
request['output'] = run(org_img_bytes, bg_img_bytes)
33+
34+
threading.Thread(target=handle_requests_by_batch).start()
35+
36+
37+
def run(org_img_bytes, bg_img_bytes):
38+
try:
39+
mask = remove(org_img_bytes, session=ort_session, only_mask=True)
40+
mask = Image.open(io.BytesIO(mask))
41+
42+
org_img = Image.open(io.BytesIO(org_img_bytes)).convert("RGB")
43+
bg_img = Image.open(io.BytesIO(bg_img_bytes)).convert("RGB")
44+
bg_img = bg_img.resize(org_img.size, Image.LANCZOS)
45+
46+
result = Image.composite(org_img, bg_img, mask)
47+
48+
bio = io.BytesIO()
49+
result.save(bio, "PNG")
50+
bio.seek(0)
51+
52+
return bio
53+
except Exception as e:
54+
return jsonify({'error': 'Exception occurs while changing background'}), 500
55+
56+
57+
@app.route("/", methods=["GET"])
58+
def index():
59+
return render_template("index.html")
60+
61+
62+
@app.route("/change-bg", methods=["POST"])
63+
def change_bg():
64+
if requests_queue.qsize() > BATCH_SIZE:
65+
return jsonify({'error': 'Too Many Reqeusts'}), 429
66+
67+
org_img_bytes = request.files['orgImage'].read()
68+
bg_img_bytes = request.files['bgImage'].read()
69+
70+
req = {
71+
'input': [org_img_bytes, bg_img_bytes]
72+
}
73+
74+
requests_queue.put(req)
75+
76+
while 'output' not in req:
77+
time.sleep(CHECK_INTERVAL)
78+
79+
io = req['output']
80+
if io == "error":
81+
return jsonify({'error': 'Server error'}), 500
82+
83+
return send_file(io, mimetype="image/png")
84+
85+
86+
@app.route("/healthz", methods=["GET"])
87+
def checkHealth():
88+
return "healthy", 200
89+
90+
91+
if __name__ == "__main__":
92+
from waitress import serve
93+
serve(app, host='0.0.0.0', port=80)

‎static/.DS_Store

6 KB
Binary file not shown.

‎static/images/.DS_Store

6 KB
Binary file not shown.

‎static/images/animal-1-out.png

2.6 MB
Loading

‎static/images/animal-1.jpg

247 KB
Loading

‎static/images/car-1-out.png

781 KB
Loading

‎static/images/car-1.jpg

166 KB
Loading

‎static/images/forest-1.jpeg

535 KB
Loading

‎static/images/sky-1.jpeg

60.2 KB
Loading

‎static/images/thumbnail-1.jpeg

8.61 KB
Loading

‎static/js/main.js

+59
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
function formSend() {
2+
const formData = new FormData();
3+
const orgImage = document.getElementById('org-img-input').files[0];
4+
const bgImage = document.getElementById('bg-img-input').files[0];
5+
6+
// original image, background image 추가하기
7+
if(orgImage == null){
8+
alert("Input the Original Image");
9+
return;
10+
}
11+
else if(bgImage == null){
12+
alert("Input the Background Image");
13+
return;
14+
}
15+
16+
formData.append("orgImage", orgImage);
17+
formData.append("bgImage", bgImage);
18+
19+
fetch(
20+
'/change-bg',
21+
{
22+
method: 'POST',
23+
body: formData,
24+
}
25+
)
26+
.then(response => {
27+
if (response.status == 200){
28+
return response
29+
}
30+
else{
31+
throw Error("Error occurs while changing background.")
32+
}
33+
})
34+
.then(response => response.blob())
35+
.then(blob => URL.createObjectURL(blob))
36+
.then(imageURL => {
37+
document.getElementById("result-img").setAttribute("src", imageURL);
38+
})
39+
.catch(e =>{
40+
})
41+
}
42+
43+
function setThumbnail(event, id){
44+
const uploaderId = id + '-input';
45+
const uploadImagePath = document.getElementById(uploaderId).value;
46+
47+
if(uploadImagePath.length == 0){
48+
document.getElementById(id).setAttribute("src", '../static/images/thumbnail-1.jpeg');
49+
document.getElementById('result-img').setAttribute("src", '../static/images/thumbnail-1.jpeg');
50+
} else{
51+
const reader = new FileReader();
52+
53+
reader.onload = function(event){
54+
document.getElementById(id).setAttribute("src", event.target.result);
55+
};
56+
57+
reader.readAsDataURL(event.target.files[0]);
58+
}
59+
}

‎templates/index.html

+88
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
<html>
2+
<head>
3+
<link rel="stylesheet" href="https://stackpath.bootstrapcdn.com/bootstrap/4.4.1/css/bootstrap.min.css" integrity="sha384-Vkoo8x4CGsO3+Hhxv8T/Q5PaXtkKtu6ug5TOeNV6gBiFeWPGFN9MuhOf23Q9Ifjh" crossorigin="anonymous">
4+
<script src="https://code.jquery.com/jquery-3.4.1.slim.min.js" integrity="sha384-J6qa4849blE2+poT4WnyKhv5vZF5SrPo0iEjwBvKU7imGFAV0wwj1yYfoRSJoZ+n" crossorigin="anonymous"></script>
5+
<script src="https://cdn.jsdelivr.net/npm/popper.js@1.16.0/dist/umd/popper.min.js" integrity="sha384-Q6E9RHvbIyZFJoft+2mJbHaEWldlvI9IOYy5n3zV9zzTtmI3UksdQRVvoxMfooAo" crossorigin="anonymous"></script>
6+
<script src="https://stackpath.bootstrapcdn.com/bootstrap/4.4.1/js/bootstrap.min.js" integrity="sha384-wfSDF2E50Y2D1uUdj0O3uMBJnjuUD4Ih7YwaYd1iqfktj0Uod8GCExl3Og8ifwB6" crossorigin="anonymous"></script>
7+
<script src="../static/js/main.js"></script>
8+
</head>
9+
<body>
10+
<div id="page-wrapper" class="container my-5">
11+
<!-- <div class="jumbotron mt-4"> -->
12+
<div id="header" class="text-center">
13+
<h1 class="mb-3">Image Background Changer</h1>
14+
<h6>Change background of image using segmentation!</h6>
15+
<a href="https://github/589hero/image-background-changer-demo" class="mr-5">GitHub</a>
16+
<a href="https://ainize.ai/589hero/image-background-changer-demo">Ainize</a>
17+
</div>
18+
<hr class="my-3">
19+
20+
<div id="explanation" class="text-center">
21+
<h2 class="mb-3">How to use</h1>
22+
<div class="text-left d-inline-flex">
23+
<p>
24+
1.&nbsp;&nbsp;Upload Original Image and Background Image.<br>
25+
2.&nbsp;&nbsp;Press "Run" Button<br>
26+
3.&nbsp;&nbsp;Check out the output image with changed background!<br>
27+
</p>
28+
</div>
29+
</div>
30+
<hr class="my-3">
31+
32+
<div id="examples" class="text-center">
33+
<h2 class="mb-3">Examples</h1>
34+
<div id="example">
35+
<figure style="display: inline-block;">
36+
<img id="sample1-1" width="250" height="200" src="../static/images/animal-1.jpg">
37+
</figure>
38+
<figure style="display: inline-block;">
39+
<img id="sample2-1" width="250" height="200" src="../static/images/forest-1.jpeg">
40+
</figure>
41+
<figure style="display: inline-block;">
42+
<img id="sample3-1" width="250" height="200" src="../static/images/animal-1-out.png">
43+
</figure>
44+
</div>
45+
<div id="example2">
46+
<figure style="display: inline-block;">
47+
<img id="sample1-2" width="250" height="200" src="../static/images/car-1.jpg">
48+
<figcaption>Original Image</figcaption>
49+
</figure>
50+
<figure style="display: inline-block;">
51+
<img id="sample2-2" width="250" height="200" src="../static/images/sky-1.jpeg">
52+
<figcaption>Background Image</figcaption>
53+
</figure>
54+
<figure style="display: inline-block;">
55+
<img id="sample3-2" width="250" height="200" src="../static/images/car-1-out.png">
56+
<figcaption>Result Image</figcaption>
57+
</figure>
58+
</div>
59+
</div>
60+
<hr class="my-3">
61+
62+
<div id="img-uploader" class="row justify-content-center">
63+
<div id="org-img-contents" class="col-4 col-6-medium col-12-small text-center">
64+
<section>
65+
<h3>Original Image</h2>
66+
<img id="org-img" class="mb-2" width="350" height="350" src="../static/images/thumbnail-1.jpeg">
67+
<input type='file' id='org-img-input' onchange="setThumbnail(event, 'org-img');">
68+
</section>
69+
</div>
70+
<div id="bg-img-contents" class="col-4 col-6-medium col-12-small text-center">
71+
<section>
72+
<h3>Background Image</h2>
73+
<img id="bg-img" class="mb-2" width="350" height="350" src="../static/images/thumbnail-1.jpeg">
74+
<input type='file' id='bg-img-input' onchange="setThumbnail(event, 'bg-img');">
75+
</section>
76+
</div>
77+
</div>
78+
<div class="my-5 text-center">
79+
<button id="run" type='submit' class="btn btn-primary" style="margin-left: auto;" onclick="formSend();"> RUN </button>
80+
</div>
81+
<div class="text-center">
82+
<h3>Result</h3>
83+
<img id="result-img" style="display: inline-block;" width="500" height="350" src="../static/images/thumbnail-1.jpeg">
84+
</div>
85+
<!-- </div> -->
86+
</div>
87+
</body>
88+
</html>

0 commit comments

Comments
 (0)
Please sign in to comment.