Farhan000's picture
Update main.py
1473cbb verified
Raw
History Blame Contribute Delete
4.06 kB
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing import image
import os
import shutil
app = FastAPI()
skin_type_model = load_model("latest_final_skin_type_model.keras")
skin_issue_model = load_model("latest_final_skin_issues_model.keras")
skin_cancer_model = load_model("latest_final_skin_cancer_model.keras")
skin_cancer_labels = ['cancer', 'no_cancer']
skin_type_labels = ['dry', 'normal', 'oily']
skin_issue_labels = ['acne', 'no_issues', 'pigmentation', 'sensitive', 'wrinkles']
UPLOAD_FOLDER = 'uploads'
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
@app.get("/")
def read_root():
return {"message": "Skin analysis API is working ✅"}
def preprocess_image(img_path, target_size=(224, 224)):
img = image.load_img(img_path, target_size=target_size)
img_array = image.img_to_array(img)
img_array = img_array / 255.0
img_array = np.expand_dims(img_array, axis=0)
return img_array
@app.post("/predict-skin-type")
async def predict_skin_type(image_file: UploadFile = File(...)):
if not image_file.filename:
raise HTTPException(status_code=400, detail="No image uploaded.")
try:
file_path = os.path.join(UPLOAD_FOLDER, image_file.filename)
with open(file_path, "wb") as buffer:
shutil.copyfileobj(image_file.file, buffer)
img_array = preprocess_image(file_path)
predictions = skin_type_model.predict(img_array)
predicted_class = skin_type_labels[np.argmax(predictions)]
confidence = float(np.max(predictions)) * 100
result = {
"predicted_class": predicted_class,
"confidence": round(confidence, 2)
}
return JSONResponse(content=result)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
finally:
if os.path.exists(file_path):
os.remove(file_path)
@app.post("/predict-skin-issue")
async def predict_skin_issue(image_file: UploadFile = File(...)):
if not image_file.filename:
raise HTTPException(status_code=400, detail="No image uploaded.")
try:
file_path = os.path.join(UPLOAD_FOLDER, image_file.filename)
with open(file_path, "wb") as buffer:
shutil.copyfileobj(image_file.file, buffer)
img_array = preprocess_image(file_path)
predictions = skin_issue_model.predict(img_array)
predicted_class = skin_issue_labels[np.argmax(predictions)]
confidence = float(np.max(predictions)) * 100
result = {
"predicted_class": predicted_class,
"confidence": round(confidence, 2)
}
return JSONResponse(content=result)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
finally:
if os.path.exists(file_path):
os.remove(file_path)
@app.post("/predict-skin-cancer")
async def predict_skin_issue(image_file: UploadFile = File(...)):
if not image_file.filename:
raise HTTPException(status_code=400, detail="No image uploaded.")
try:
file_path = os.path.join(UPLOAD_FOLDER, image_file.filename)
with open(file_path, "wb") as buffer:
shutil.copyfileobj(image_file.file, buffer)
img_array = preprocess_image(file_path)
predictions = skin_cancer_model.predict(img_array)
predicted_class = skin_cancer_labels[np.argmax(predictions)]
confidence = float(np.max(predictions)) * 100
result = {
"predicted_class": predicted_class,
"confidence": round(confidence, 2)
}
return JSONResponse(content=result)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
finally:
if os.path.exists(file_path):
os.remove(file_path)
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)