app.py

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import FileResponse
import cv2
from PIL import Image
import re
import gdown
import os
import uuid

URL = "https://leekwoon-edge-api.hf.space"
# URL = "http://localhost:7860"

app = FastAPI(
    version="0.0.1",
    servers=[
        {
            "url": URL,
            "description": "image edge detection API",
        }
    ],
)


def extract_file_id(drive_url: str) -> str:
    """
    Google Drive URL에서 파일 ID를 추출합니다.

    Parameters:
    drive_url (str): Google Drive 파일 URL.

    Returns:
    str: 추출된 파일 ID.
    """
    # 정규 표현식을 사용하여 URL에서 파일 ID 추출
    match = re.search(r'/d/([a-zA-Z0-9_-]+)', drive_url)
    if match:
        return match.group(1)
    
    # 다른 URL 형식에서 파일 ID 추출
    match = re.search(r'file/d/([a-zA-Z0-9_-]+)', drive_url)
    if match:
        return match.group(1)
    
    # 공유 링크에서 파일 ID 추출
    match = re.search(r'([a-zA-Z0-9_-]{33,})', drive_url)
    if match:
        return match.group(1)

    raise ValueError("Invalid Google Drive URL")


@app.post("/detect-edges/")
def detect_edges(url: str, request: Request):
    try:
        # Google Drive 파일 ID 추출
        file_id = extract_file_id(url)
        download_url = f"https://drive.google.com/uc?id={file_id}"
        
        # 이미지를 Google Drive에서 다운로드
        temp_input_image = f'/tmp/{uuid.uuid4()}.png'
        gdown.download(download_url, temp_input_image, quiet=False)

        # 이미지를 읽어들임
        image = cv2.imread(temp_input_image)

        # 이미지가 제대로 읽혔는지 확인
        if image is None:
            raise HTTPException(status_code=400, detail="Invalid image URL or file not found.")

        # 그레이스케일로 변환
        gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

        # Canny 에지 검출 적용
        edges = cv2.Canny(gray_image, threshold1=100, threshold2=200)
        edges_image = Image.fromarray(edges)

        # 흰색 배경 생성
        white_background = Image.new("RGB", edges_image.size, (255, 255, 255))

        # Canny 에지를 흰색 배경 위에 합성
        edges_on_white = Image.composite(Image.new("RGB", edges_image.size, (0, 0, 0)), white_background, edges_image)

        # Sobel 연산자 적용
        sobelx = cv2.Sobel(gray_image, cv2.CV_64F, 1, 0, ksize=5)  # X 축에 대한 Sobel 연산자
        sobely = cv2.Sobel(gray_image, cv2.CV_64F, 0, 1, ksize=5)  # Y 축에 대한 Sobel 연산자

        # 두 기울기 결합
        sobel_combined = cv2.magnitude(sobelx, sobely)

        # 결과를 [0, 255] 범위로 정규화
        sobel_combined = cv2.normalize(sobel_combined, None, 0, 255, cv2.NORM_MINMAX)
        sobel_combined = sobel_combined.astype('uint8')

        # 결과를 PIL 이미지로 변환
        sobel_image = Image.fromarray(sobel_combined)

        # Sobel 에지를 흰색 배경 위에 합성
        edges_on_white_sobel = Image.composite(Image.new("RGB", sobel_image.size, (0, 0, 0)), white_background, sobel_image)

        # 최종 Sobel 이미지를 파일로 저장 (PNG 형식)
        output_file_path = f'/tmp/{uuid.uuid4()}.png'
        edges_on_white_sobel.save(output_file_path, format='PNG')

        # 입력 이미지 임시 파일 삭제
        os.remove(temp_input_image)

        # 서버의 URL을 기반으로 다운로드 URL 생성
        download_url = os.path.join(URL, "tmp", os.path.basename(output_file_path))

        # 클라이언트에게 결과 이미지 파일 경로를 반환
        return {"download_url": download_url}

    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.get("/tmp/{filename}")
def download_file(filename: str):
    file_path = f"/tmp/{filename}"
    if os.path.exists(file_path):
        return FileResponse(file_path, media_type='image/png', filename=filename)
    else:
        raise HTTPException(status_code=404, detail="File not found")