diff --git a/main.py b/main.py new file mode 100644 index 0000000..6ba13f7 --- /dev/null +++ b/main.py @@ -0,0 +1,62 @@ +from fastapi import FastAPI, File, UploadFile +from fastapi.responses import JSONResponse +from transformers import pipeline +from PIL import Image +from cachetools import Cache +import io, hashlib, logging + +app = FastAPI() + +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') + +# Initialize Cache with no TTL +cache = Cache(maxsize=1000) + +# Load the model using the transformers pipeline +model = pipeline("image-classification", model="falconsai/nsfw_image_detection") + +def hash_data(data): + return hashlib.sha256(data).hexdigest() + +@app.post("/api/v1/detect") +async def classify_image(file: UploadFile = File(...)): + try: + logging.info(f"Processing {file.filename}") + # Read the image file + image_data = await file.read() + image_hash = hash_data(image_data) + + if image_hash in cache: + # Return cached entry + logging.info(f"Returning cached entry for {file.filename}") + return JSONResponse(status_code=200, content=cache[image_hash]) + + image = Image.open(io.BytesIO(image_data)) + + # Use the model to classify the image + results = model(image) + + # Find the prediction with the highest confidence using the max() function + best_prediction = max(results, key=lambda x: x['score']) + + # Calculate the confidence score, rounded to the nearest tenth and as a percentage + confidence_percentage = round(best_prediction['score'] * 100, 1) + + # Prepare the custom response data + response_data = { + "filename": file.filename, + "isNsfw": best_prediction['label'] == 'nsfw', + "confidence_percentage": confidence_percentage + } + + # Populate hash + cache[image_hash] = response_data + + return JSONResponse(status_code=200, content=response_data) + + except Exception as e: + return JSONResponse(status_code=500, content={"message": str(e)}) + +if __name__ == "__main__": + import uvicorn + uvicorn.run(app, host="127.0.0.1", port=8000) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..c26f0a8 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,9 @@ +fastapi==0.110.2 +uvicorn[standard]==0.29.0 +transformers==4.40.0 +aiohttp==3.9.5 +pillow==10.3.0 +python-multipart==0.0.9 +tensorflow==2.16.1 +tf-keras==2.16.0 +cachetools===5.3.3