import base64 import hashlib import json import random import uuid import cv2 from flask import Flask, request, jsonify import sys import os from PIL import Image import io import numpy as np import websocket import openpose_gen as opg from comfy_socket import get_images from postprocessing import expo_shuffle_image_steps, expo_add_to_background_image, expo_postprocess_main import skeleton_lib as skel sys.path.append('./') app = Flask(__name__) info = json.load(open('info.json')) comfyui_address = info['comfyui_address'] expo_raw_sd_dir = info['expo_raw_sd_dir'] expo_openpose_dir = info['expo_openpose_dir'] expo_postprocessed_dir = info['expo_postprocessed_dir'] expo_postprocess_temp_dir = info['expo_postprocess_temp_dir'] on_postprocessing = False on_testing = True @app.route('/expo_fencing_pose', methods=['POST']) def expo_fencing_pose(): if on_postprocessing: return jsonify({"status": "error", "message": "Postprocessing in progress"}), 503 if request.is_json: data = request.get_json() coordinates = data['coordinates'] canvas_size = data['canvas_size'] batch = data['batch'] step = data['step'] if coordinates is None or canvas_size is None or 'batch' not in data or 'step' not in data: return jsonify({"status": "error", "message": "Missing data"}), 422 right_fencer_coordinates = get_predicted_coordinates(coordinates) left_fencer_dir = os.path.join(expo_openpose_dir, 'left_fencer') os.makedirs(left_fencer_dir, exist_ok=True) right_fencer_dir = os.path.join(expo_openpose_dir, 'right_fencer') os.makedirs(right_fencer_dir, exist_ok=True) left_openpose_image_path = opg.expo_save_bodypose(canvas_size[0], canvas_size[1], coordinates, batch, step, left_fencer_dir, skel.coco_limbSeq, skel.coco_colors) right_openpose_image_path = opg.expo_save_bodypose(canvas_size[0], canvas_size[1], right_fencer_coordinates, batch, step, right_fencer_dir, skel.coco_limbSeq, skel.coco_colors) left_fencer_raw_image_dir = os.path.join(expo_raw_sd_dir, 'left_fencer') os.makedirs(left_fencer_raw_image_dir, exist_ok=True) right_fencer_raw_image_dir = os.path.join(expo_raw_sd_dir, 'right_fencer') os.makedirs(right_fencer_raw_image_dir, exist_ok=True) expo_fencer_prompt(left_openpose_image_path, left_fencer_raw_image_dir, batch, step) expo_fencer_prompt(right_openpose_image_path, right_fencer_raw_image_dir, batch, step) return jsonify({"status": "success", "message": "Data received"}), 201 else: return jsonify({"status": "error", "message": "Request must be JSON"}), 415 def get_predicted_coordinates(coordinates: list, width: int, height: int) -> list: # TODO implement the model to predict the right fencer's coordinates # coordinates = [x1, y1, c1, x2, y2, c2, ...], # where x, y are the coordinates and c is the confidence score # there should be 18 keypoints from 0 to 17 # they are not normalized, they are by the size of the width and height # the the limbSeq and colors of points need to convert from and to skel.coco_limbSeq, skel.coco_colors # those are in skeleton_lib.py # when testing, can visualize with the method expo_save_bodypose in openpose_gen.py return mirror_coordinates(coordinates, width) def mirror_coordinates(coordinates: list, width: int) -> list: for i in range(0, len(coordinates), 3): coordinates[i] = width - coordinates[i] return coordinates def expo_fencer_prompt(openpose_image_path, save_dir, batch, step): prompt = json.loads(open("./prompts/fencer_03.json", "r", encoding="utf-8").read()) openpose_image_name = opg.upload_image(openpose_image_path) opg.upload_image("./images/ref_black.png", "ref_black.png") print(openpose_image_name) prompt["3"]["inputs"]["seed"] = random.randint(0, 10000000000) prompt["29"]["inputs"]['image'] = "ref_black.png" prompt["17"]["inputs"]['image'] = openpose_image_name client_id = hashlib.sha256(str(random.getrandbits(256)).encode('utf-8')).hexdigest() ws = websocket.WebSocket() ws.connect("ws://{}/ws?clientId={}".format(comfyui_address, client_id)) images = get_images(ws, prompt, client_id) for node_id in images: for idx, image_data in enumerate(images[node_id]): image = Image.open(io.BytesIO(image_data)) image_path = os.path.join(save_dir, f"{batch}_{step}.png") image.save(image_path) def expo_clear_images(): if on_testing: return for root, dirs, files in os.walk(expo_openpose_dir): for file in files: os.remove(os.path.join(root, file)) for root, dirs, files in os.walk(expo_raw_sd_dir): for file in files: os.remove(os.path.join(root, file)) for root, dirs, files in os.walk(expo_postprocess_temp_dir): for file in files: os.remove(os.path.join(root, file)) @app.route('/expo_postprocess', methods=['POST']) def expo_postprocess(): if on_postprocessing: return jsonify({"status": "error", "message": "Postprocessing in progress"}), 503 print("Postprocessing") on_postprocessing = True expo_postprocess_main() expo_clear_images() on_postprocessing = False @app.route('/gen_image', methods=['POST']) def gen_image(): if request.is_json: data = request.get_json() coordinates = data['coordinates'] canvas_size = data['canvas_size'] pid = data['pid'] if not coordinates or not canvas_size: return jsonify({"status": "error", "message": "Missing data"}), 422 openpose_image_path = opg.save_bodypose(canvas_size[0], canvas_size[1], coordinates, pid) # gen_fencer_prompt(openpose_image_path, pid, comfyui_address) return jsonify({"status": "success", "message": "Data received"}), 201 else: return jsonify({"status": "error", "message": "Request must be JSON"}), 415 @app.route('/gen_group_pic', methods=['POST']) def gen_group_pic(): if request.is_json: data = request.get_json() coordinates_list = data['coordinates_list'] canvas_size = data['canvas_size'] pid = data['pid'] base_image = base64.b64decode(data['base_image']) if not coordinates_list or not canvas_size or not base_image or not pid: return jsonify({"status": "error", "message": "Missing data"}), 422 for i in range(len(coordinates_list)): coordinates_list[i] = coordinates_list[i]['coordinates'] openpose_image_path = opg.save_bodypose_mulit(canvas_size[0], canvas_size[1], coordinates_list, pid) gen_group_pic_prompt(openpose_image_path, base_image, pid, comfyui_address) return jsonify({"status": "success", "message": "Data received"}), 201 else: return jsonify({"status": "error", "message": "Request must be JSON"}), 415 def gen_fencer_prompt(openpose_image_path, pid, comfyUI_address): with open("./prompts/fencerAPI.json", "r") as f: prompt_json = f.read() prompt = json.loads(prompt_json) openpose_image_name = opg.upload_image_circular_queue(openpose_image_path, 20, pid, comfyUI_address) opg.upload_image("./images/ref_black.png", "ref_black.png") prompt["3"]["inputs"]["seed"] = random.randint(0, 10000000000) prompt["29"]["inputs"]['image'] = "./images/ref_black.png" prompt["17"]["inputs"]['image'] = openpose_image_name opg.queue_prompt(prompt, comfyUI_address) def gen_group_pic_prompt(openpose_image_path, base_image, pid, comfyUI_address): with open("./prompts/group_pic.json", "r") as f: prompt_json = f.read() prompt = json.loads(prompt_json) openpose_image_name = opg.upload_image_circular_queue(openpose_image_path, 30, pid, comfyUI_address) base_image_name = opg.upload_image_circular_queue(base_image, 30, pid, comfyUI_address) prompt["3"]["inputs"]["seed"] = random.randint(0, 10000000000) prompt["10"]["inputs"]['image'] = openpose_image_name prompt["14"]["inputs"]['image'] = base_image_name opg.queue_prompt(prompt, comfyUI_address) if __name__ == '__main__': expo_postprocess() # app.run(debug=True)