Spaces:
Sleeping
Sleeping
| from fastapi import FastAPI, UploadFile, File | |
| import tensorflow as tf | |
| import numpy as np | |
| from PIL import Image | |
| import io | |
| import requests | |
| import os | |
| MODEL_URL = "https://huggingface.co/blaxx14/cat-vs-dog-inceptionv3/resolve/main/cat_dog_inception_v3.h5" | |
| MODEL_PATH = "cat_dog_inception_v3.h5" | |
| app = FastAPI() | |
| def download_model(): | |
| if not os.path.exists(MODEL_PATH): | |
| print("Downloading model...") | |
| response = requests.get(MODEL_URL) | |
| with open(MODEL_PATH, "wb") as f: | |
| f.write(response.content) | |
| print("Model downloaded!") | |
| download_model() | |
| print("Loading model...") | |
| model = tf.keras.models.load_model(MODEL_PATH) | |
| print("Model loaded!") | |
| def preprocess_image(image): | |
| img = Image.open(io.BytesIO(image)).convert("RGB") | |
| img = img.resize((150, 150)) | |
| img = np.array(img) / 255.0 | |
| img = np.expand_dims(img, axis=0) | |
| return img | |
| async def predict(file: UploadFile = File(...)): | |
| image = await file.read() | |
| processed_img = preprocess_image(image) | |
| prediction = model.predict(processed_img) | |
| return {"prediction": prediction.tolist()} | |