Update app.py

app.py

@@ -1,151 +1,334 @@
-import hashlib
-import os
-from io import BytesIO
-
-import gradio as gr
-import grpc
-from PIL import Image
-from cachetools import LRUCache
-
-from inference_pb2 import HairSwapRequest, HairSwapResponse
-from inference_pb2_grpc import HairSwapServiceStub
-from utils.shape_predictor import align_face
-
-
-def get_bytes(img):
-    if img is None:
-        return img
-
-    buffered = BytesIO()
-    img.save(buffered, format="JPEG")
-    return buffered.getvalue()
-
-
-def bytes_to_image(image: bytes) -> Image.Image:
-    image = Image.open(BytesIO(image))
-    return image
-
-
-def center_crop(img):
-    width, height = img.size
-    side = min(width, height)
-
-    left = (width - side) / 2
-    top = (height - side) / 2
-    right = (width + side) / 2
-    bottom = (height + side) / 2
-
-    img = img.crop((left, top, right, bottom))
-    return img
-
-
-def resize(name):
-    def resize_inner(img, align):
-        global align_cache
-
-        if name in align:
-            img_hash = hashlib.md5(get_bytes(img)).hexdigest()
-
-            if img_hash not in align_cache:
-                img = align_face(img, return_tensors=False)[0]
-                align_cache[img_hash] = img
-            else:
-                img = align_cache[img_hash]
-
-        elif img.size != (1024, 1024):
-            img = center_crop(img)
-            img = img.resize((1024, 1024), Image.Resampling.LANCZOS)
-
-        return img
-
-    return resize_inner
-
-
-def swap_hair(face, shape, color, blending, poisson_iters, poisson_erosion):
-    if not face and not shape and not color:
-        return gr.update(visible=False), gr.update(value="Need to upload a face and at least a shape or color ❗", visible=True)
-    elif not face:
-        return gr.update(visible=False), gr.update(value="Need to upload a face ❗", visible=True)
-    elif not shape and not color:
-        return gr.update(visible=False), gr.update(value="Need to upload at least a shape or color ❗", visible=True)
-
-    face_bytes, shape_bytes, color_bytes = map(lambda item: get_bytes(item), (face, shape, color))
-
-    if shape_bytes is None:
-        shape_bytes = b'face'
-    if color_bytes is None:
-        color_bytes = b'shape'
-
-    with grpc.insecure_channel(os.environ['SERVER']) as channel:
-        stub = HairSwapServiceStub(channel)
-
-        output: HairSwapResponse = stub.swap(
-            HairSwapRequest(face=face_bytes, shape=shape_bytes, color=color_bytes, blending=blending,
-                            poisson_iters=poisson_iters, poisson_erosion=poisson_erosion, use_cache=True)
-        )
-
-    output = bytes_to_image(output.image)
-    return gr.update(value=output, visible=True), gr.update(visible=False)
-
-
-def get_demo():
-    with gr.Blocks() as demo:
-        gr.Markdown("## HairFastGan")
-        gr.Markdown(
-            '<div style="display: flex; align-items: center; gap: 10px;">'
-            '<span>Official HairFastGAN Gradio demo:</span>'
-            '<a href="https://arxiv.org/abs/2404.01094"><img src="https://img.shields.io/badge/arXiv-2404.01094-b31b1b.svg" height=22.5></a>'
-            '<a href="https://github.com/AIRI-Institute/HairFastGAN"><img src="https://img.shields.io/badge/github-%23121011.svg?style=for-the-badge&logo=github&logoColor=white" height=22.5></a>'
-            '<a href="https://huggingface.co/AIRI-Institute/HairFastGAN"><img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/model-on-hf-md.svg" height=22.5></a>'
-            '<a href="https://colab.research.google.com/#fileId=https://huggingface.co/AIRI-Institute/HairFastGAN/blob/main/notebooks/HairFast_inference.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" height=22.5></a>'
-            '</div>'
-        )
-        with gr.Row():
-            with gr.Column():
-                source = gr.Image(label="Source photo to try on the hairstyle", type="pil")
-                with gr.Row():
-                    shape = gr.Image(label="Shape photo with desired hairstyle (optional)", type="pil")
-                    color = gr.Image(label="Color photo with desired hair color (optional)", type="pil")
-                with gr.Accordion("Advanced Options", open=False):
-                    blending = gr.Radio(["Article", "Alternative_v1", "Alternative_v2"], value='Article',
-                                        label="Color Encoder version", info="Selects a model for hair color transfer.")
-                    poisson_iters = gr.Slider(0, 2500, value=0, step=1, label="Poisson iters",
-                                              info="The power of blending with the original image, helps to recover more details. Not included in the article, disabled by default.")
-                    poisson_erosion = gr.Slider(1, 100, value=15, step=1, label="Poisson erosion",
-                                                info="Smooths out the blending area.")
-                    align = gr.CheckboxGroup(["Face", "Shape", "Color"], value=["Face", "Shape", "Color"],
-                                             label="Image cropping [recommended]",
-                                             info="Selects which images to crop by face")
-                btn = gr.Button("Get the haircut")
-            with gr.Column():
-                output = gr.Image(label="Your result")
-                error_message = gr.Textbox(label="⚠️ Error ⚠️", visible=False, elem_classes="error-message")
-
-        gr.Examples(examples=[["input/0.png", "input/1.png", "input/2.png"], ["input/6.png", "input/7.png", None],
-                              ["input/10.jpg", None, "input/11.jpg"]],
-                    inputs=[source, shape, color], outputs=output)
-
-        source.upload(fn=resize('Face'), inputs=[source, align], outputs=source)
-        shape.upload(fn=resize('Shape'), inputs=[shape, align], outputs=shape)
-        color.upload(fn=resize('Color'), inputs=[color, align], outputs=color)
-
-        btn.click(fn=swap_hair, inputs=[source, shape, color, blending, poisson_iters, poisson_erosion],
-                  outputs=[output, error_message])
-
-        gr.Markdown('''To cite the paper by the authors
-    ```
-        @article{nikolaev2024hairfastgan,
-          title={HairFastGAN: Realistic and Robust Hair Transfer with a Fast Encoder-Based Approach},
-          author={Nikolaev, Maxim and Kuznetsov, Mikhail and Vetrov, Dmitry and Alanov, Aibek},
-          journal={arXiv preprint arXiv:2404.01094},
-          year={2024}
-        }
-    ```
-        ''')
-    return demo
-
-
-if __name__ == '__main__':
-    align_cache = LRUCache(maxsize=10)
-    demo = get_demo()
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+import hashlib
+import os
+import base64
+from io import BytesIO
+import threading
+
+import gradio as gr
+from flask import Flask, request, jsonify
+from flask_cors import CORS
+from PIL import Image
+from cachetools import LRUCache
+import grpc
+
+from inference_pb2 import HairSwapRequest, HairSwapResponse
+from inference_pb2_grpc import HairSwapServiceStub
+from utils.shape_predictor import align_face
+
+# Initialize Flask app
+flask_app = Flask(__name__)
+CORS(flask_app)
+
+# Global cache
+align_cache = LRUCache(maxsize=10)
+
+
+def get_bytes(img):
+    if img is None:
+        return img
+    buffered = BytesIO()
+    img.save(buffered, format="JPEG")
+    return buffered.getvalue()
+
+
+def bytes_to_image(image: bytes) -> Image.Image:
+    image = Image.open(BytesIO(image))
+    return image
+
+
+def base64_to_image(base64_string: str) -> Image.Image:
+    """Convert base64 string to PIL Image"""
+    image_data = base64.b64decode(base64_string.split(',')[-1])
+    return Image.open(BytesIO(image_data))
+
+
+def image_to_base64(img: Image.Image) -> str:
+    """Convert PIL Image to base64 string"""
+    buffered = BytesIO()
+    img.save(buffered, format="JPEG")
+    img_str = base64.b64encode(buffered.getvalue()).decode()
+    return f"data:image/jpeg;base64,{img_str}"
+
+
+def center_crop(img):
+    width, height = img.size
+    side = min(width, height)
+    left = (width - side) / 2
+    top = (height - side) / 2
+    right = (width + side) / 2
+    bottom = (height + side) / 2
+    img = img.crop((left, top, right, bottom))
+    return img
+
+
+def resize(name):
+    def resize_inner(img, align):
+        global align_cache
+
+        if name in align:
+            img_hash = hashlib.md5(get_bytes(img)).hexdigest()
+
+            if img_hash not in align_cache:
+                img = align_face(img, return_tensors=False)[0]
+                align_cache[img_hash] = img
+            else:
+                img = align_cache[img_hash]
+
+        elif img.size != (1024, 1024):
+            img = center_crop(img)
+            img = img.resize((1024, 1024), Image.Resampling.LANCZOS)
+
+        return img
+
+    return resize_inner
+
+
+def swap_hair_core(face, shape, color, blending, poisson_iters, poisson_erosion, align_settings):
+    """Core hair swap function used by both Gradio and Flask"""
+    if not face:
+        return None, "Need to upload a face ❗"
+    
+    if not shape and not color:
+        return None, "Need to upload at least a shape or color ❗"
+
+    # Process images
+    face_img = face
+    shape_img = shape
+    color_img = color
+
+    # Apply alignment if needed
+    if 'Face' in align_settings:
+        img_hash = hashlib.md5(get_bytes(face_img)).hexdigest()
+        if img_hash not in align_cache:
+            face_img = align_face(face_img, return_tensors=False)[0]
+            align_cache[img_hash] = face_img
+        else:
+            face_img = align_cache[img_hash]
+    
+    if shape_img and 'Shape' in align_settings:
+        img_hash = hashlib.md5(get_bytes(shape_img)).hexdigest()
+        if img_hash not in align_cache:
+            shape_img = align_face(shape_img, return_tensors=False)[0]
+            align_cache[img_hash] = shape_img
+        else:
+            shape_img = align_cache[img_hash]
+    
+    if color_img and 'Color' in align_settings:
+        img_hash = hashlib.md5(get_bytes(color_img)).hexdigest()
+        if img_hash not in align_cache:
+            color_img = align_face(color_img, return_tensors=False)[0]
+            align_cache[img_hash] = color_img
+        else:
+            color_img = align_cache[img_hash]
+
+    # Resize if needed
+    if face_img.size != (1024, 1024):
+        face_img = center_crop(face_img)
+        face_img = face_img.resize((1024, 1024), Image.Resampling.LANCZOS)
+    
+    if shape_img and shape_img.size != (1024, 1024):
+        shape_img = center_crop(shape_img)
+        shape_img = shape_img.resize((1024, 1024), Image.Resampling.LANCZOS)
+    
+    if color_img and color_img.size != (1024, 1024):
+        color_img = center_crop(color_img)
+        color_img = color_img.resize((1024, 1024), Image.Resampling.LANCZOS)
+
+    # Convert to bytes
+    face_bytes = get_bytes(face_img)
+    shape_bytes = get_bytes(shape_img) if shape_img else b'face'
+    color_bytes = get_bytes(color_img) if color_img else b'shape'
+
+    try:
+        with grpc.insecure_channel(os.environ.get('SERVER', 'localhost:50051')) as channel:
+            stub = HairSwapServiceStub(channel)
+            output: HairSwapResponse = stub.swap(
+                HairSwapRequest(
+                    face=face_bytes,
+                    shape=shape_bytes,
+                    color=color_bytes,
+                    blending=blending,
+                    poisson_iters=poisson_iters,
+                    poisson_erosion=poisson_erosion,
+                    use_cache=True
+                )
+            )
+        
+        output_img = bytes_to_image(output.image)
+        return output_img, None
+    
+    except Exception as e:
+        return None, f"Error: {str(e)}"
+
+
+def swap_hair(face, shape, color, blending, poisson_iters, poisson_erosion, align):
+    """Gradio interface function"""
+    result, error = swap_hair_core(
+        face, shape, color, blending, 
+        poisson_iters, poisson_erosion, align
+    )
+    
+    if error:
+        return gr.update(visible=False), gr.update(value=error, visible=True)
+    
+    return gr.update(value=result, visible=True), gr.update(visible=False)
+
+
+# Flask API Endpoints
+@flask_app.route('/health', methods=['GET'])
+def health_check():
+    """Health check endpoint"""
+    return jsonify({"status": "healthy", "service": "HairFastGAN API"}), 200
+
+
+@flask_app.route('/api/swap-hair', methods=['POST'])
+def api_swap_hair():
+    """Flask API endpoint for hair swap"""
+    try:
+        data = request.get_json()
+        
+        if not data:
+            return jsonify({"error": "No JSON data provided"}), 400
+        
+        if 'face' not in data:
+            return jsonify({"error": "Face image is required"}), 400
+        
+        if 'shape' not in data and 'color' not in data:
+            return jsonify({"error": "At least shape or color image is required"}), 400
+        
+        # Parse images
+        face_img = base64_to_image(data['face'])
+        shape_img = base64_to_image(data['shape']) if 'shape' in data and data['shape'] else None
+        color_img = base64_to_image(data['color']) if 'color' in data and data['color'] else None
+        
+        # Get options
+        blending = data.get('blending', 'Article')
+        poisson_iters = int(data.get('poisson_iters', 0))
+        poisson_erosion = int(data.get('poisson_erosion', 15))
+        
+        # Build align settings
+        align_settings = []
+        if data.get('align_face', True):
+            align_settings.append('Face')
+        if data.get('align_shape', True):
+            align_settings.append('Shape')
+        if data.get('align_color', True):
+            align_settings.append('Color')
+        
+        # Process
+        result_img, error = swap_hair_core(
+            face_img, shape_img, color_img,
+            blending, poisson_iters, poisson_erosion,
+            align_settings
+        )
+        
+        if error:
+            return jsonify({"error": error}), 500
+        
+        # Convert to base64
+        result_base64 = image_to_base64(result_img)
+        
+        return jsonify({
+            "success": True,
+            "result": result_base64,
+            "message": "Hair swap completed successfully"
+        }), 200
+        
+    except Exception as e:
+        return jsonify({"error": f"Internal server error: {str(e)}"}), 500
+
+
+@flask_app.route('/', methods=['GET'])
+def api_root():
+    """API documentation"""
+    return jsonify({
+        "service": "HairFastGAN API",
+        "version": "1.0",
+        "endpoints": {
+            "/health": "GET - Health check",
+            "/api/swap-hair": "POST - Hair swap endpoint"
+        }
+    }), 200
+
+
+# Gradio Interface
+def get_demo():
+    with gr.Blocks() as demo:
+        gr.Markdown("## HairFastGan")
+        gr.Markdown(
+            '<div style="display: flex; align-items: center; gap: 10px;">'
+            '<span>Official HairFastGAN Gradio demo:</span>'
+            '<a href="https://arxiv.org/abs/2404.01094"><img src="https://img.shields.io/badge/arXiv-2404.01094-b31b1b.svg" height=22.5></a>'
+            '<a href="https://github.com/AIRI-Institute/HairFastGAN"><img src="https://img.shields.io/badge/github-%23121011.svg?style=for-the-badge&logo=github&logoColor=white" height=22.5></a>'
+            '<a href="https://huggingface.co/AIRI-Institute/HairFastGAN"><img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/model-on-hf-md.svg" height=22.5></a>'
+            '</div>'
+        )
+        
+        gr.Markdown("### 🔗 API Endpoint Available!")
+        gr.Markdown("Use `/api/swap-hair` endpoint for programmatic access. See API docs at `/` endpoint.")
+        
+        with gr.Row():
+            with gr.Column():
+                source = gr.Image(label="Source photo to try on the hairstyle", type="pil")
+                with gr.Row():
+                    shape = gr.Image(label="Shape photo with desired hairstyle (optional)", type="pil")
+                    color = gr.Image(label="Color photo with desired hair color (optional)", type="pil")
+                with gr.Accordion("Advanced Options", open=False):
+                    blending = gr.Radio(["Article", "Alternative_v1", "Alternative_v2"], value='Article',
+                                        label="Color Encoder version", info="Selects a model for hair color transfer.")
+                    poisson_iters = gr.Slider(0, 2500, value=0, step=1, label="Poisson iters",
+                                              info="The power of blending with the original image, helps to recover more details. Not included in the article, disabled by default.")
+                    poisson_erosion = gr.Slider(1, 100, value=15, step=1, label="Poisson erosion",
+                                                info="Smooths out the blending area.")
+                    align = gr.CheckboxGroup(["Face", "Shape", "Color"], value=["Face", "Shape", "Color"],
+                                             label="Image cropping [recommended]",
+                                             info="Selects which images to crop by face")
+                btn = gr.Button("Get the haircut")
+            with gr.Column():
+                output = gr.Image(label="Your result")
+                error_message = gr.Textbox(label="⚠️ Error ⚠️", visible=False, elem_classes="error-message")
+
+        gr.Examples(examples=[["input/0.png", "input/1.png", "input/2.png"], ["input/6.png", "input/7.png", None],
+                              ["input/10.jpg", None, "input/11.jpg"]],
+                    inputs=[source, shape, color], outputs=output)
+
+        source.upload(fn=resize('Face'), inputs=[source, align], outputs=source)
+        shape.upload(fn=resize('Shape'), inputs=[shape, align], outputs=shape)
+        color.upload(fn=resize('Color'), inputs=[color, align], outputs=color)
+
+        btn.click(fn=swap_hair, inputs=[source, shape, color, blending, poisson_iters, poisson_erosion, align],
+                  outputs=[output, error_message])
+
+        gr.Markdown('''To cite the paper by the authors
+    ```
+        @article{nikolaev2024hairfastgan,
+          title={HairFastGAN: Realistic and Robust Hair Transfer with a Fast Encoder-Based Approach},
+          author={Nikolaev, Maxim and Kuznetsov, Mikhail and Vetrov, Dmitry and Alanov, Aibek},
+          journal={arXiv preprint arXiv:2404.01094},
+          year={2024}
+        }
+    ```
+        ''')
+    return demo
+
+
+def run_flask():
+    """Run Flask in a separate thread"""
+    flask_app.run(host='0.0.0.0', port=5000, debug=False, use_reloader=False)
+
+
+if __name__ == '__main__':
+    # Start Flask API in background thread
+    flask_thread = threading.Thread(target=run_flask, daemon=True)
+    flask_thread.start()
+    
+    print("🚀 Flask API running on http://0.0.0.0:5000")
+    print("🎨 Gradio UI starting on http://0.0.0.0:7860")
+    
+    # Start Gradio
+    demo = get_demo()
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )