Upload 6 files

Dockerfile

@@ -0,0 +1,25 @@
+# Use an official Python runtime as a parent image
+FROM python:3.12-slim
+
+# Set the working directory in the container
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    gcc \
+    build-essential \
+    pkg-config \
+    libhdf5-dev \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy the current directory contents into the container at /app
+COPY . /app
+
+# Install any needed packages specified in requirements.txt
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Make port 8000 available to the world outside this container
+EXPOSE 8000
+
+# Command to run the Uvicorn server
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Amir Boroumand
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

SECURITY.md

@@ -0,0 +1,10 @@
+# Security Policy
+
+## Supported Versions
+
+This is an open source project that is provided as-is without warranty or liability.
+As such no supportability commitment. The maintainers will do the best they can to address any report promptly and responsibly.
+
+## Reporting a Vulnerability
+
+Please use the "Private vulnerability reporting" feature in the GitHub repository (under the "Security" tab).

main.py

@@ -0,0 +1,174 @@
+"""Module providing an API for NSFW image detection."""
+
+import io
+import hashlib
+import logging
+import aiohttp
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import JSONResponse
+from transformers import pipeline
+from transformers.pipelines import PipelineException
+from PIL import Image
+from cachetools import Cache
+import tensorflow as tf
+from models import (
+    FileImageDetectionResponse,
+    UrlImageDetectionResponse,
+    ImageUrlsRequest,
+)
+
+
+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")
+
+# Detect the device used by TensorFlow
+DEVICE = "GPU" if tf.config.list_physical_devices("GPU") else "CPU"
+logging.info("TensorFlow version: %s", tf.__version__)
+logging.info("Model is using: %s", DEVICE)
+
+if DEVICE == "GPU":
+    logging.info("GPUs available: %d", len(tf.config.list_physical_devices("GPU")))
+
+
+async def download_image(image_url: str) -> bytes:
+    """Download an image from a URL."""
+    async with aiohttp.ClientSession() as session:
+        async with session.get(image_url) as response:
+            if response.status != 200:
+                raise HTTPException(
+                    status_code=response.status, detail="Image could not be retrieved."
+                )
+            return await response.read()
+
+
+def hash_data(data):
+    """Function for hashing image data."""
+    return hashlib.sha256(data).hexdigest()
+
+
+@app.post("/v1/detect", response_model=FileImageDetectionResponse)
+async def classify_image(file: UploadFile = File(None)):
+    """Function analyzing image."""
+    if file is None:
+        raise HTTPException(
+            status_code=400,
+            detail="An image file must be provided.",
+        )
+
+    try:
+        logging.info("Processing %s", 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("Returning cached entry for %s", file.filename)
+
+            cached_response = cache[image_hash]
+            response_data = {**cached_response, "file_name": file.filename}
+
+            return FileImageDetectionResponse(**response_data)
+
+        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 = {
+            "is_nsfw": best_prediction["label"] == "nsfw",
+            "confidence_percentage": confidence_percentage,
+        }
+
+        # Populate hash
+        cache[image_hash] = response_data.copy()
+
+        # Add file_name to the API response
+        response_data["file_name"] = file.filename
+
+        return FileImageDetectionResponse(**response_data)
+
+    except PipelineException as e:
+        logging.error("Error processing image: %s", str(e))
+        raise HTTPException(
+            status_code=500, detail=f"Error processing image: {str(e)}"
+        ) from e
+
+
+@app.post("/v1/detect/urls", response_model=list[UrlImageDetectionResponse])
+async def classify_images(request: ImageUrlsRequest):
+    """Function analyzing images from URLs."""
+    response_data = []
+
+    for image_url in request.urls:
+        try:
+            logging.info("Downloading image from URL: %s", image_url)
+            image_data = await download_image(image_url)
+            image_hash = hash_data(image_data)
+
+            if image_hash in cache:
+                # Return cached entry
+                logging.info("Returning cached entry for %s", image_url)
+
+                cached_response = cache[image_hash]
+                response = {**cached_response, "url": image_url}
+
+                response_data.append(response)
+                continue
+
+            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
+            detection_result = {
+                "is_nsfw": best_prediction["label"] == "nsfw",
+                "confidence_percentage": confidence_percentage,
+            }
+
+            # Populate hash
+            cache[image_hash] = detection_result.copy()
+
+            # Add url to the API response
+            detection_result["url"] = image_url
+
+            response_data.append(detection_result)
+
+        except PipelineException as e:
+            logging.error("Error processing image from %s: %s", image_url, str(e))
+            raise HTTPException(
+                status_code=500,
+                detail=f"Error processing image from {image_url}: {str(e)}",
+            ) from e
+
+    return JSONResponse(status_code=200, content=response_data)
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host="127.0.0.1", port=8000)

models.py

@@ -0,0 +1,49 @@
+"""Module providing base models."""
+
+from pydantic import BaseModel
+
+
+class ImageUrlsRequest(BaseModel):
+    """
+    Model representing the request body for the /v1/detect/urls endpoint.
+
+    Attributes:
+        urls (list[str]): List of image URLs to be processed.
+    """
+
+    urls: list[str]
+
+
+class ImageDetectionResponse(BaseModel):
+    """
+    Base model representing the response body for image detection.
+
+    Attributes:
+        is_nsfw (bool): Whether the image is classified as NSFW.
+        confidence_percentage (float): Confidence level of the NSFW classification.
+    """
+
+    is_nsfw: bool
+    confidence_percentage: float
+
+
+class FileImageDetectionResponse(ImageDetectionResponse):
+    """
+    Model extending ImageDetectionResponse with a file attribute.
+
+    Attributes:
+        file (str): The name of the file that was processed.
+    """
+
+    file_name: str
+
+
+class UrlImageDetectionResponse(ImageDetectionResponse):
+    """
+    Model extending ImageDetectionResponse with a URL attribute.
+
+    Attributes:
+        url (str): The URL of the image that was processed.
+    """
+
+    url: str

requirements.txt

@@ -0,0 +1,10 @@
+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
+pydantic===2.7.2