app.py

from flask import Flask, request, send_file
import cv2
import numpy as np
import tensorflow as tf
from io import BytesIO
from PIL import Image

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

app = Flask(__name__)

# Ping API Route
@app.route('/ping', methods=['GET'])
def ping():
    return 'Ping from Python API!'

def send_image_response(image, filename="image.jpg"):
    """Helper function to return image response"""
    _, img_encoded = cv2.imencode('.jpg', image)
    img_bytes = img_encoded.tobytes()
    buffer = BytesIO(img_bytes)
    buffer.seek(0)

    # Set Cache-Control headers to prevent caching
    response = send_file(buffer, mimetype='image/jpeg', as_attachment=True, download_name=filename)
    response.cache_control.no_cache = True
    response.cache_control.no_store = True
    response.cache_control.must_revalidate = True
    response.headers['Pragma'] = 'no-cache'
    response.headers['Expires'] = '0'

    return response

@app.before_request
def before_request():
    """Reset or cleanup any global state before processing each request."""
    # Any necessary cleanup can be done here (e.g., reset global variables).
    pass

@app.after_request
def after_request(response):
    """Perform cleanup after each request if necessary."""
    # Post-request cleanup, if needed, can be done here (e.g., free resources).
    return response

# Pencil Sketch API Route
@app.route('/sketch', methods=['POST'])
def sketch_image():
    if 'image' not in request.files:
        return "No image part", 400
    file = request.files['image']
    img_array = np.frombuffer(file.read(), np.uint8)
    image = cv2.imdecode(img_array, cv2.IMREAD_COLOR)

    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    inverted_image = 255 - gray_image
    blurred_image = cv2.GaussianBlur(inverted_image, (111, 111), 0)
    inverted_blurred = 255 - blurred_image
    sketch = cv2.divide(gray_image, inverted_blurred, scale=256.0)

    return send_image_response(sketch, "sketch.jpg")

# Sharpened Pencil Sketch API Route
@app.route('/sharpened_sketch', methods=['POST'])
def sharpened_sketch_image():
    if 'image' not in request.files:
        return "No image part", 400
    file = request.files['image']
    img_array = np.frombuffer(file.read(), np.uint8)
    image = cv2.imdecode(img_array, cv2.IMREAD_COLOR)

    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    inverted_image = 255 - gray_image
    blurred_image = cv2.GaussianBlur(inverted_image, (21, 21), 0)
    inverted_blurred = 255 - blurred_image
    sketch = cv2.divide(gray_image, inverted_blurred, scale=256.0)

    sharpen_kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])
    sharpened_sketch = cv2.filter2D(sketch, -1, sharpen_kernel)

    return send_image_response(sharpened_sketch, "sharpened_sketch.jpg")

# Anime-style effect API Route
@app.route('/anime_effect', methods=['POST'])
def anime_effect_image():
    if 'image' not in request.files:
        return "No image part", 400
    file = request.files['image']
    img_array = np.frombuffer(file.read(), np.uint8)
    image = cv2.imdecode(img_array, cv2.IMREAD_COLOR)

    bilateral_filtered = cv2.bilateralFilter(image, d=9, sigmaColor=300, sigmaSpace=300)

    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    gray = cv2.medianBlur(gray, 5)

    edges = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 9, 10)

    data = np.float32(image).reshape((-1, 3))
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
    k = 8
    _, labels, centers = cv2.kmeans(data, k, None, criteria, 10, cv2.KMEANS_RANDOM_CENTERS)
    centers = np.uint8(centers)
    quantized_image = centers[labels.flatten()]
    quantized_image = quantized_image.reshape(image.shape)

    cartoon = cv2.bitwise_and(bilateral_filtered, bilateral_filtered, mask=edges)
    anime_image = cv2.bitwise_and(quantized_image, cartoon)

    return send_image_response(anime_image, "anime_effect.jpg")

# Cartoon Effect API Route
@app.route('/cartoon_effect', methods=['POST'])
def cartoon_effect_image():
    if 'image' not in request.files:
        return "No image part", 400
    file = request.files['image']
    img_array = np.frombuffer(file.read(), np.uint8)
    image = cv2.imdecode(img_array, cv2.IMREAD_COLOR)

    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    bilateral_image = cv2.bilateralFilter(image, d=9, sigmaColor=75, sigmaSpace=75)

    edges = cv2.adaptiveThreshold(gray_image, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 9, 9)
    edges_colored = cv2.cvtColor(edges, cv2.COLOR_GRAY2BGR)

    cartoon_image = cv2.bitwise_and(bilateral_image, edges_colored)

    return send_image_response(cartoon_image, "cartoon_effect.jpg")

# Cartoon Conversion API Route
@app.route('/cartoon_convert', methods=['POST'])
def cartoon_convert():
    if 'image' not in request.files:
        return "No image part", 400
    
    file = request.files['image']
    img_array = np.frombuffer(file.read(), np.uint8)
    image = cv2.imdecode(img_array, cv2.IMREAD_COLOR)

    model_path = '/app/1.tflite'

    source_img = load_image_from_bytes(img_array)
    if source_img is None:
        return "Error loading image", 400

    processed_img = preprocess_image(source_img)

    try:
        interpreter = tf.lite.Interpreter(model_path=model_path)
        input_details = interpreter.get_input_details()
        interpreter.allocate_tensors()
        interpreter.set_tensor(input_details[0]['index'], processed_img)
        interpreter.invoke()
        result = interpreter.tensor(interpreter.get_output_details()[0]['index'])()
    except Exception as e:
        return f"Error during model inference: {e}", 500

    output = post_process(result, preserve_details=True)

    output_image = Image.fromarray(output)
    img_bytes = BytesIO()
    output_image.save(img_bytes, format='JPEG')
    img_bytes.seek(0)

    # Set Cache-Control headers to prevent caching
    response = send_file(img_bytes, mimetype='image/jpeg', as_attachment=True, download_name='cartoon_converted_image.jpg')
    response.cache_control.no_cache = True
    response.cache_control.no_store = True
    response.cache_control.must_revalidate = True
    response.headers['Pragma'] = 'no-cache'
    response.headers['Expires'] = '0'

    return response

# Helper Functions for Cartoon Conversion
def load_image_from_bytes(img_bytes):
    try:
        img = cv2.imdecode(np.frombuffer(img_bytes, np.uint8), cv2.IMREAD_COLOR)
        if img is None:
            raise ValueError("Failed to load image")
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = img.astype(np.float32) / 127.5 - 1
        img = np.expand_dims(img, 0)
        return tf.convert_to_tensor(img)
    except Exception as e:
        print(f"Error loading image: {e}")
        return None

def preprocess_image(img, target_dim=512):
    shape = tf.cast(tf.shape(img)[1:-1], tf.float32)
    min_dim = tf.reduce_min(shape)
    scale = target_dim / min_dim
    new_shape = tf.cast(shape * scale, tf.int32)

    img = tf.image.resize(img, new_shape, method='area')
    img = tf.image.resize_with_crop_or_pad(img, target_dim, target_dim)

    return img

def apply_detail_preserving_smoothing(img):
    smooth = cv2.detailEnhance(img, sigma_s=10, sigma_r=0.15)
    smooth = cv2.edgePreservingFilter(smooth, flags=1, sigma_s=60, sigma_r=0.4)
    return smooth

def enhance_colors(img):
    lab = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)
    l, a, b = cv2.split(lab)

    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    l = clahe.apply(l)

    lab = cv2.merge((l,a,b))
    enhanced = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)
    return enhanced

def post_process(output, preserve_details=True):
    img = (np.squeeze(output) + 1.0) * 127.5
    img = np.clip(img, 0, 255).astype(np.uint8)

    if preserve_details:
        img = apply_detail_preserving_smoothing(img)
        img = enhance_colors(img)
        img = cv2.bilateralFilter(img, 5, 35, 35)

    return img


if __name__ == '__main__':
    app.run(debug=True, host="0.0.0.0", port=7860)