File size: 3,486 Bytes
9adfe8b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
from io import BytesIO, StringIO
import cv2
from anime_face_detector import create_detector
from werkzeug.wsgi import FileWrapper
from flask import Flask, request, Response, send_file
import math
import numpy as np
def get_deg(arr):
rad = math.atan2(arr[3]-arr[1],arr[2]-arr[0])
PI = math.pi
deg = (rad*180)/PI
return deg
detector = create_detector('yolov3', device='cpu')
gg = cv2.imread('gg.png', cv2.IMREAD_UNCHANGED)
gg = cv2.cvtColor(gg, cv2.COLOR_BGRA2RGBA)
def generate(image_file: BytesIO) -> bytes:
encoded = np.asarray(bytearray(image_file.read()), dtype=np.uint8)
img = cv2.imdecode(encoded, cv2.IMREAD_COLOR)
preds = detector(img)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGBA)
# for face in preds:
# draw.rectangle((face['bbox'][0], face['bbox'][1], face['bbox'][2], face['bbox'][3]), outline=(255, 0, 0), width=5)
# x = face['bbox'][0]
# y = face['bbox'][1]
# for i, point in enumerate(face['keypoints']):
# # draw.ellipse((point[0]-2, point[1]-2, point[0]+2, point[1]+2), fill=(255, 0, 0))
# draw.text((point[0], point[1]), str(i), font=ImageFont.truetype('arial.ttf', 10), fill=(255, 0, 0))
if len(preds) == 0:
return False
for face in preds:
points = face['keypoints']
color = img[int(points[27][1]), int(points[27][0])+10]
polygon = np.array([
[points[0][0], points[0][1]],
[points[1][0], points[1][1]],
[points[2][0], points[2][1]],
[points[3][0], points[4][1]],
[points[4][0], points[4][1]],
[points[10][0], points[10][1]],
[points[9][0], points[9][1]],
[points[8][0], points[8][1]],
[points[7][0], points[7][1]],
[points[6][0], points[6][1]],
[points[5][0], points[5][1]]
], np.int32)
cv2.fillConvexPoly(img, polygon, color=(int(color[0]), int(color[1]), int(color[2]), 255))
deg = get_deg([points[0][0], points[0][1], points[4][0], points[4][1]])
rotated = gg.copy()
resize = math.sqrt((points[10][0] - points[5][0])**2 + (points[10][1] - points[5][1])**2)
rotated = cv2.resize(rotated, (int(resize), int(resize*1.12)))
matrix = cv2.getPerspectiveTransform(
np.float32([[0, 0], [rotated.shape[0],0], [0, rotated.shape[1]], [rotated.shape[0],rotated.shape[1]]]),
np.float32([[points[5][0], points[5][1]], [points[10][0], points[10][1]], [points[1][0], points[1][1]], [points[3][0], points[3][1]]]))
rotated = cv2.warpPerspective(rotated, matrix, (img.shape[1], img.shape[0]))
alpha = rotated[:, :, 3] / 255.
for i in range(3):
pointx, pointy = points[5][:2]
pointx, pointy = int(pointx), int(pointy)
img[:, :, i] = (1. - alpha) * img[0:, 0:, i] + alpha * rotated[:, :, i]
buffer = cv2.imencode('.png', cv2.cvtColor(img, cv2.COLOR_RGBA2BGRA))[1]
return buffer.tobytes()
app = Flask(__name__)
@app.post('/generate')
def index():
if request.files.get('file') is None:
return "no file", 400
file = request.files.get('file')
dst = BytesIO()
file.save(dst)
dst.seek(0)
result = generate(dst)
if not result:
return {"status": 400}, 400
return Response(result, mimetype='image/png', direct_passthrough=True)
if __name__ == "__main__":
app.run("0.0.0.0", 8080, debug=False) |