34fd2wqc/MachineLearning

1from model import build_dce_net2from PIL import Image3import numpy as np4import tensorflow as tf5import io6import base647from fastapi import FastAPI, File, UploadFile8from fastapi.responses import JSONResponse, HTMLResponse9 10app = FastAPI()11 12# Load model13model = build_dce_net()14model.load_weights("best_model.h5")15 16def enhance_image(image: Image.Image) -> Image.Image:17    width, height = image.size18    new_width = min(width // 2, 500)19    new_height = min(height // 2, 500)20    img = image.resize((new_width, new_height))21    img_array = np.array(img).astype("float32") / 255.022    img_array = np.expand_dims(img_array, axis=0)23 24    output = model(img_array)25    x = img_array26 27    for i in range(8):28        ri = output[:, :, :, i * 3:(i + 1) * 3]29        x = x + 2.0 * ri * (tf.square(x) - x)30 31    result = tf.clip_by_value(x[0], 0.0, 1.0).numpy()32 33    brightness_scale = 0.6 / (result.mean() + 1e-6)34    result *= brightness_scale35    result = np.clip(result, 0.0, 1.0)36 37    result_image = Image.fromarray((result * 255).astype("uint8"))38 39    return result_image40 41@app.get("/")42def home():43    html = """44    <h3>Zero-DCE API is running</h3>45    <p><a href='/enhance'>Go to Enhancement Page</a></p>46    """47    return HTMLResponse(content=html)48 49@app.get("/enhance")50def form():51    html = """52    <form action="/enhance" enctype="multipart/form-data" method="post">53        <input name="file" type="file" accept="image/*">54        <input type="submit" value="Enhance">55    </form>56    """57    return HTMLResponse(content=html)58 59@app.post("/enhance")60async def enhance(file: UploadFile = File(...)):61    try:62        contents = await file.read()63        image = Image.open(io.BytesIO(contents)).convert("RGB")64        enhanced = enhance_image(image)65 66        # Encode image to base6467        buf = io.BytesIO()68        enhanced.save(buf, format="PNG")69        encoded = base64.b64encode(buf.getvalue()).decode("utf-8")70 71        return JSONResponse(content={"result": encoded})72    except Exception as e:73        return JSONResponse(status_code=500, content={"error": str(e)})