import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
from tensorflow.keras.applications.vgg16 import preprocess_input

class Desk_classifier:
    def __init__(self, threshold = 0.5):
        #threshold, model_path 추후 수정
        model_path = '...'
        self.model = load_model(model_path)
        self.threshold = threshold

    def predict(self, img_path: str) -> bool:
        img = image.load_img(img_path, target_size = (224, 224))
        x = image.img_to_array(img)
        x = np.expand_dims(x, axis = 0)
        x = preprocess_input(x)

        prob = self.model.predict(x)[0][0]
        return bool(prob >= self.threshold)

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from desk_classify import Desk_classifier
import requests
import shutil
import os
import uuid

app = FastAPI()
classifier = Desk_classifier()

class ImageRequest(BaseModel):
    image_url: str

@app.post("/classify")
async def classify_image_url(request: ImageRequest):
    temp_dir = "temp_downloads"
    os.makedirs(temp_dir, exist_ok=True)

    temp_path = os.path.join(temp_dir, f"{uuid.uuid4().hex}.jpg")

    try:
        response = requests.get(request.image_url, stream=True)
        if response.status_code != 200:
            raise HTTPException(status_code=400, detail="Failed to download image from URL")

        with open(temp_path, "wb") as f:
            for chunk in response.iter_content(1024):
                f.write(chunk)

        result = classifier.predict(temp_path)
        return {"is_desk": result}

    finally:
        if os.path.exists(temp_path):
            os.remove(temp_path)