main.py

import io
import tempfile
from functools import lru_cache

import boto3
from config import Settings
from fastapi import FastAPI, UploadFile

from app import color


@lru_cache
def get_settings():
    return Settings()


settings = get_settings()

s3 = boto3.client(
    "s3",
    endpoint_url=settings.endpoint_url,
    aws_access_key_id=settings.aws_access_key_id,
    aws_secret_access_key=settings.aws_secret_access_key,
)


def old_photo_restoration(file):
    with tempfile.NamedTemporaryFile() as fp:
        fp.write(file.file.read())
        result = color(fp.name)

        if len(result) < 2:
            return None

        file_path = upload_to_s3(result[1])
        fp.close()
        return file_path


def upload_to_s3(file_path):
    with open(file_path, "rb") as file:
        file_content = file.read()

        s3.upload_fileobj(io.BytesIO(file_content), "old-photo-restoration", file.name)
        return f"{settings.cloud_flare_endpoint}/{file.name}"


app = FastAPI()


@app.get("/")
def read_root():
    return {"Hello": "World"}


@app.post("/upload")
def upload_image(file: UploadFile):
    return old_photo_restoration(file)