img2isometric

 

 

1. main.py

import uvicorn
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import FileResponse, HTMLResponse
from fastapi.staticfiles import StaticFiles
from rembg import remove
from PIL import Image
import io
import os
import datetime
import uuid # 고유 파일명 생성을 위해 추가
import sqlite3
import logging # 로깅 추가

# --- 설정 ---
app = FastAPI()

# 환경 변수 사용 (또는 기본값)
UPLOAD_DIR = os.getenv("UPLOAD_DIR", "uploads")
PROCESSED_DIR = os.getenv("PROCESSED_DIR", "processed")
DB_NAME = os.getenv("DB_NAME", "image_db.sqlite")
MAX_FILE_SIZE_MB = int(os.getenv("MAX_FILE_SIZE_MB", "10")) # 최대 파일 크기 10MB
ALLOWED_EXTENSIONS = {"png", "jpg", "jpeg", "gif", "webp"} # 허용된 확장자

# 로깅 설정
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 업로드/처리된 이미지를 저장할 폴더 생성
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(PROCESSED_DIR, exist_ok=True)

# --- DB 설정 ---
def init_db():
    conn = None
    try:
        conn = sqlite3.connect(DB_NAME)
        cursor = conn.cursor()
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS images (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            original_filename TEXT NOT NULL, # 원본 업로드 시의 파일명 (클라이언트 기준)
            original_server_path TEXT NOT NULL, # 서버에 저장된 원본 파일의 경로
            processed_server_path TEXT NOT NULL, # 서버에 저장된 처리된 파일의 경로
            processed_url TEXT NOT NULL, # 클라이언트가 접근할 수 있는 처리된 이미지 URL
            uploaded_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
        ''')
        conn.commit()
        logger.info(f"Database {DB_NAME} initialized.")
    except sqlite3.Error as e:
        logger.error(f"Error initializing database: {e}")
    finally:
        if conn:
            conn.close()

init_db()

# --- 헬퍼 함수 ---
def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

# 이미지 파일의 매직 넘버(시그니처) 검사 (일부 주요 이미지 형식만)
def is_valid_image(file_content):
    if file_content.startswith(b'\x89PNG\r\n\x1a\n'): # PNG
        return True
    if file_content.startswith(b'\xFF\xD8\xFF'): # JPEG
        return True
    if file_content.startswith(b'GIF87a') or file_content.startswith(b'GIF89a'): # GIF
        return True
    if file_content.startswith(b'RIFF') and b'WEBP' in file_content: # WEBP (간단 검사)
        return True
    return False

# --- API 엔드포인트 ---
app.mount("/static", StaticFiles(directory="static"), name="static")
app.mount("/processed_images", StaticFiles(directory=PROCESSED_DIR), name="processed_images") # URL 경로 변경

@app.post("/process-image/")
async def process_image_endpoint(file: UploadFile = File(...)):
    if not file.filename:
        raise HTTPException(status_code=400, detail="파일 이름이 없습니다.")

    # 1. 파일 확장자 검사
    if not allowed_file(file.filename):
        logger.warning(f"업로드 거부: 허용되지 않는 확장자 - {file.filename}")
        raise HTTPException(status_code=400, detail="허용되지 않는 이미지 파일 형식입니다. (png, jpg, jpeg, gif, webp만 가능)")

    # 2. 파일 크기 제한
    file_content = await file.read()
    if len(file_content) > MAX_FILE_SIZE_MB * 1024 * 1024:
        logger.warning(f"업로드 거부: 파일 크기 초과 - {file.filename} ({len(file_content) / (1024*1024):.2f} MB)")
        raise HTTPException(status_code=413, detail=f"파일 크기는 {MAX_FILE_SIZE_MB}MB를 초과할 수 없습니다.")
    
    # 3. 파일 시그니처 검사
    if not is_valid_image(file_content):
        logger.warning(f"업로드 거부: 유효하지 않은 이미지 시그니처 - {file.filename}")
        raise HTTPException(status_code=400, detail="유효하지 않은 이미지 파일입니다.")

    # 4. 고유한 파일 이름 생성 (UUID 사용)
    file_extension = file.filename.rsplit('.', 1)[1].lower()
    unique_id = str(uuid.uuid4())
    original_server_filename = f"{unique_id}_original.{file_extension}"
    processed_server_filename = f"{unique_id}_processed.png" # rembg는 PNG로 출력하므로 고정

    original_server_path = os.path.join(UPLOAD_DIR, original_server_filename)
    processed_server_path = os.path.join(PROCESSED_DIR, processed_server_filename)
    
    # 클라이언트가 접근할 수 있는 URL (StaticFiles 마운트 경로와 일치)
    processed_client_url = f"/processed_images/{processed_server_filename}"

    # 5. 원본 이미지 저장
    try:
        with open(original_server_path, "wb") as f:
            f.write(file_content)
        logger.info(f"원본 이미지 저장: {original_server_path}")
    except IOError as e:
        logger.error(f"원본 이미지 저장 오류: {e}")
        raise HTTPException(status_code=500, detail="원본 이미지 저장에 실패했습니다.")

    # 6. 이미지 배경 제거
    try:
        output_bytes = remove(file_content)
        output_image = Image.open(io.BytesIO(output_bytes))
        
        # 7. 결과 이미지 저장
        output_image.save(processed_server_path)
        logger.info(f"처리된 이미지 저장: {processed_server_path}")

        # 8. 데이터베이스에 파일 정보 저장
        conn = None
        try:
            conn = sqlite3.connect(DB_NAME)
            cursor = conn.cursor()
            cursor.execute(
                "INSERT INTO images (original_filename, original_server_path, processed_server_path, processed_url) VALUES (?, ?, ?, ?)",
                (file.filename, original_server_path, processed_server_path, processed_client_url)
            )
            conn.commit()
            logger.info(f"DB에 이미지 정보 저장 완료: {processed_client_url}")
        except sqlite3.Error as db_e:
            logger.error(f"DB 저장 오류: {db_e}")
            raise HTTPException(status_code=500, detail="데이터베이스 저장에 실패했습니다.")
        finally:
            if conn:
                conn.close()

        # 9. 클라이언트에 결과 반환
        return {
            "message": "처리 완료",
            "original_url": original_server_path, # 개발 편의상 제공, 실제 서비스에서는 보안상 내부 경로 노출 지양
            "processed_url": processed_client_url
        }
        
    except Exception as e:
        logger.error(f"이미지 처리 중 오류 발생: {e}", exc_info=True)
        # 이미 원본 파일이 저장되었을 수 있으므로 정리 (선택 사항, 복잡해질 수 있음)
        # if os.path.exists(original_server_path): os.remove(original_server_path)
        raise HTTPException(status_code=500, detail="이미지 처리 중 예상치 못한 오류가 발생했습니다.")

@app.get("/", response_class=HTMLResponse)
async def read_index():
    with open('static/index.html', 'r', encoding='utf-8') as f:
        return HTMLResponse(content=f.read())

# --- 서버 실행 ---
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)