new stuff

.dockerignore

@@ -0,0 +1,4 @@
+.venv
+.vscode
+data
+__pycache__

.gitignore

@@ -0,0 +1,4 @@
+.venv
+.vscode
+data
+__pycache__

Dockerfile

@@ -0,0 +1,15 @@
+FROM python:3.9-slim
+
+WORKDIR /code
+
+COPY ./requirements.txt /code/requirements.txt
+
+RUN apt-get update && apt-get install ffmpeg libsm6 libxext6 -y
+
+RUN pip install torch==1.9.0+cpu torchvision==0.10.0+cpu -f https://download.pytorch.org/whl/torch_stable.html
+
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+
+COPY . /code/
+
+CMD ["uvicorn", "main:app",  "--host", "0.0.0.0", "--port", "8000"]

main.py

@@ -0,0 +1,95 @@
+# Import Fast API
+from fastapi import FastAPI, Request, UploadFile, File
+from fastapi.templating import Jinja2Templates
+from fastapi.responses import StreamingResponse
+
+# Import bytes
+from io import BytesIO
+import os
+
+# Import logging
+import logging
+
+# Import utilities
+from src.utils.utils import IMAGE_FORMATS
+
+# Import machine learning
+from src.predict import predict
+from ultralytics import YOLO
+from huggingface_hub import hf_hub_download
+
+
+# Initialazing FastAPI application
+app = FastAPI()
+
+# Initialazing templates
+templates = Jinja2Templates(directory="templates")
+
+# Initialazing logger
+logger = logging.getLogger(__name__)
+
+
+logger.info(f"Loading YOLO model...")
+
+# Download YOLO model from Hugging Face Hub
+model_path = hf_hub_download(
+    repo_id="arnabdhar/YOLOv8-Face-Detection", filename="model.pt"
+)
+
+
+# Load YOLO model
+model = YOLO(model_path)
+
+
+# Index route
+@app.get("/")
+async def root(request: Request):
+    context = {"request": request}
+    # Render index.html
+    return templates.TemplateResponse("index.html", context)
+
+
+# Upload images decorator
+@app.post("/predict-img")
+def predict_image(file: UploadFile = File(...)):
+    try:
+        # Try to read the file
+        contents = file.file.read()
+
+        # Open file and write contents
+        with open(file.filename, "wb") as f:
+            f.write(contents)
+
+        # Get image filename
+        image = file.filename
+
+        # Check if image format is valid
+        if not image.endswith(IMAGE_FORMATS):
+            # If not, raise an error
+            raise ValueError("Invalid image format")
+
+    except Exception as e:
+        # If there is an error, return the error
+        return {f"{e}"}
+
+    finally:
+        file.file.close()
+
+    # Getting image path
+    image = file.filename
+
+    # Predicting
+    results = predict(model, image)
+
+    # TODO
+    # extract extension from image and use it to save the image
+    # Convert image to bytes
+    img_bytes = BytesIO()
+    results.save(img_bytes, "JPEG")
+    img_bytes.seek(0)
+
+    # Removing the image
+    os.remove(image)
+
+    # Render image
+    return StreamingResponse(content=img_bytes, media_type="image/jpeg")

requirements.txt

@@ -0,0 +1,6 @@
+fastapi
+pillow
+ultralytics
+huggingface-hub
+
+

src/predict.py

@@ -0,0 +1,80 @@
+# Import machine learning
+from ultralytics import YOLO
+
+# Import computer vision
+import cv2
+from PIL import Image
+
+
+def predict(detection_model: YOLO, image_path: str) -> Image:
+    """
+    The function predicts with the given model and returns the image with the bounding boxes
+
+    Args:
+    detection_model (YOLO): The YOLO model
+    image_path (str): The path to the image
+
+    Returns:
+    Image: The image with the bounding boxes
+    """
+    # Read the image
+    img = cv2.imread(image_path)
+
+    # Predict with the model
+    output = detection_model.predict(image_path, verbose=False)
+
+    # If output is not None then draw the bounding boxes
+    if output:
+        # Loop through the output
+        for bbox in output:
+            for i, p in enumerate(bbox):
+
+                # Converting the p to cpu
+                p = p.to("cpu")
+
+                # Conveting to numpy
+                p = p.numpy()
+
+                # Extracting the coords
+                coords = p.boxes.xyxy[0]
+
+                # Extracting the coords
+                xmin = coords[0]
+                ymin = coords[1]
+                xmax = coords[2]
+                ymax = coords[3]
+
+                # Converting to int
+                xmin = int(xmin)
+                ymin = int(ymin)
+                xmax = int(xmax)
+                ymax = int(ymax)
+
+                # Extracting the prob
+                prob = p.boxes.conf[0]
+
+                # Drawing the bounding box
+                cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)
+
+                # Drawing the text with the probability
+                cv2.putText(
+                    img,
+                    f"{prob:.2f}",
+                    (xmin, ymin - 10),
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                    0.9,
+                    (36, 255, 12),
+                    2,
+                )
+
+        # Converting the image to RGB
+        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+        # Converting the image to PIL Image object
+        pil_image = Image.fromarray(img_rgb)
+
+        # Return the image
+        return pil_image
+
+    # If output is None then return the original image
+    return Image.open(image_path)

src/utils/utils.py

@@ -0,0 +1,16 @@
+IMAGE_FORMATS = (
+    "jpg",
+    "jpeg",
+    "png",
+    "bmp",
+    "gif",
+    "tiff",
+    "webp",
+    "JPG",
+    "JPEG",
+    "PNG",
+    "BMP",
+    "GIF",
+    "TIFF",
+    "WEBP",
+)

templates/index.html

@@ -0,0 +1,61 @@
+<!DOCTYPE html>
+<html lang="en">
+  <head>
+    <!-- Required meta tags -->
+    <meta charset="utf-8" />
+    <meta name="viewport" content="width=device-width, initial-scale=1" />
+    <meta
+      name="viewport"
+      content="width=device-width, initial-scale=1, shrink-to-fit=no"
+    />
+
+    <!-- Bootstrap CSS -->
+    <link
+      href="https://cdn.jsdelivr.net/npm/[email protected]/dist/css/bootstrap.min.css"
+      rel="stylesheet"
+      integrity="sha384-EVSTQN3/azprG1Anm3QDgpJLIm9Nao0Yz1ztcQTwFspd3yD65VohhpuuCOmLASjC"
+      crossorigin="anonymous"
+    />
+
+    <script
+      src="https://cdn.jsdelivr.net/npm/[email protected]/dist/js/bootstrap.bundle.min.js"
+      integrity="sha384-MrcW6ZMFYlzcLA8Nl+NtUVF0sA7MsXsP1UyJoMp4YLEuNSfAP+JcXn/tWtIaxVXM"
+      crossorigin="anonymous"
+    ></script>
+
+    <title>Hello, world!</title>
+  </head>
+  <body>
+    <div class="w-25 p-4 mx-auto">
+    <form method="post" action="/predict-img" enctype="multipart/form-data">
+      
+        <input
+          class="form-control"
+          type="file"
+          id="formFileMultiple"
+          multiple
+          name="file"
+        />
+        <button type="submit" class="btn btn-primary">
+          Upload files
+        </button>
+      </div>
+    </form>
+    <!-- Optional JavaScript; choose one of the two! -->
+
+    <!-- Option 1: Bootstrap Bundle with Popper -->
+
+    <!-- Option 2: Separate Popper and Bootstrap JS -->
+
+    <script
+      src="https://cdn.jsdelivr.net/npm/@popperjs/[email protected]/dist/umd/popper.min.js"
+      integrity="sha384-IQsoLXl5PILFhosVNubq5LC7Qb9DXgDA9i+tQ8Zj3iwWAwPtgFTxbJ8NT4GN1R8p"
+      crossorigin="anonymous"
+    ></script>
+    <script
+      src="https://cdn.jsdelivr.net/npm/[email protected]/dist/js/bootstrap.min.js"
+      integrity="sha384-cVKIPhGWiC2Al4u+LWgxfKTRIcfu0JTxR+EQDz/bgldoEyl4H0zUF0QKbrJ0EcQF"
+      crossorigin="anonymous"
+    ></script>
+  </body>
+</html>