add face recognition (#10)

README.md

@@ -38,6 +38,17 @@ serve.run_server()
 
 or, run `python -m fastserve.models --model sdxl-turbo --batch_size 2 --timeout 1` from terminal.
 
+### Face  Detection
+
+```python
+from fastserve.models import FaceDetection
+
+serve = FaceDetection(batch_size=2, timeout=1)
+serve.run_server()
+```
+
+or, run `python -m fastserve.models --model face-detection --batch_size 2 --timeout 1` from terminal.
+
 ### Serve Custom Model
 
 To serve a custom model, you will have to implement `handle` method for `FastServe` that processes a batch of inputs and

src/fastserve/models/__init__.py

@@ -1,3 +1,4 @@
+from fastserve.models.face_reco import FaceDetection as FaceDetection
 from fastserve.models.llama_cpp import ServeLlamaCpp as ServeLlamaCpp
 from fastserve.models.sdxl_turbo import ServeSDXLTurbo as ServeSDXLTurbo
 from fastserve.models.ssd import ServeSSD1B as ServeSSD1B

src/fastserve/models/__main__.py

@@ -1,5 +1,6 @@
 import argparse
 
+from fastserve.models.face_reco import FaceDetection
 from fastserve.models.llama_cpp import ServeLlamaCpp
 from fastserve.models.sdxl_turbo import ServeSDXLTurbo
 from fastserve.models.ssd import ServeSSD1B
@@ -49,6 +50,12 @@
         timeout=args.timeout,
         batch_size=args.batch_size,
     )
+elif args.model == "face-detection":
+    app = FaceDetection(
+        timeout=args.timeout,
+        batch_size=args.batch_size,
+    )
+
 else:
     raise Exception(f"FastServe.models doesn't implement model={args.model}")
 

src/fastserve/models/face_reco.py

@@ -0,0 +1,33 @@
+import logging
+from typing import List
+
+from fastapi import UploadFile
+
+from fastserve import FastServe
+
+logging.getLogger(__name__)
+
+try:
+    import face_recognition
+except ModuleNotFoundError:
+    face_recognition = None
+
+
+class FaceDetection(FastServe):
+    def __init__(self, batch_size=1, timeout=0):
+        super().__init__(
+            batch_size=batch_size, timeout=timeout, input_schema=UploadFile
+        )
+        if not face_recognition:
+            raise ModuleNotFoundError(
+                "face_recognition is not installed. "
+                "Please run 'pip install face-recognition' first."
+            )
+
+    def handle(self, batch: List[UploadFile]) -> List[List[float]]:
+        results = []
+        for file in batch:
+            image = face_recognition.load_image_file(file.file)
+            face_locations = face_recognition.face_locations(image)
+            results.append(face_locations)
+        return results