add face detection model

main.py

@@ -1,15 +1,20 @@
-from typing import Union
+from typing import Optional, Union
 import dotenv
 import traceback
 import json
 import io
 import os
 import base64
-from fastapi import FastAPI, File, HTTPException, UploadFile, Response
+from fastapi import FastAPI, File, Form, HTTPException, UploadFile, Response
+from ultralytics import YOLO
 import models.face_classifier as classifier
 from fastapi.middleware.cors import CORSMiddleware
+from huggingface_hub import hf_hub_download
 from PIL import Image
 from rembg import remove
+import procedures
+from utils import handlers
+from supervision import Detections
 from utils.helpers import combine_images, image_to_base64, calculate_mask_area, process_image
 
 
@@ -31,7 +36,7 @@ app.add_middleware(
 
 model = classifier.FaceSegmentationModel()
 
-
+yolo_model_path = hf_hub_download(repo_id="arnabdhar/YOLOv8-Face-Detection", filename="model.pt")
 
 @app.post("/segment/", summary="Classify skin type based on image given",tags=["Classify"])
 async def predict_image(file: UploadFile = File(...)):
@@ -91,4 +96,49 @@ async def predict_image(file: UploadFile = File(...)):
     except Exception as e:
         # Mendapatkan stack trace
         error_traceback = traceback.format_exc()
-        raise HTTPException(status_code=500, detail=f"An error occurred: {error_traceback}")
+        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")
+
+@app.post("/detect-face/", summary="Detect face from image", tags=["Classify"])
+async def detect_face(file: Optional[UploadFile] = File(None), base64_string: Optional[str] = Form(None)):
+    try:
+        if file is None and base64_string is None:
+            raise HTTPException(status_code=400, detail="No input data provided")
+        base64_handler = handlers.Base64Handler()
+        image_handler = handlers.ImageFileHandler(successor=base64_handler)
+        input_data: Union[UploadFile, str, None] = file if file is not None else base64_string
+        print(input_data)
+        pil_image = await image_handler.handle(input_data)
+        if pil_image is None:
+            raise HTTPException(status_code=400, detail="Unsupported file type")
+        
+        # Load the YOLO model
+
+        model = YOLO(yolo_model_path)
+        
+        # Inference using the pil image
+        output = model(pil_image)
+        results = Detections.from_ultralytics(output[0])
+        
+        if len(results) == 0:
+            raise HTTPException(status_code=404, detail="No face detected")
+        
+        # Get the first bounding box
+        first_bbox = results[0].xyxy[0].tolist()
+        
+        # Crop the image using the bounding box
+        x_min, y_min, x_max, y_max = map(int, first_bbox)
+        cropped_image = pil_image.crop((x_min, y_min, x_max, y_max))
+        
+        # Convert cropped image to Base64
+        buffered = io.BytesIO()
+        cropped_image.save(buffered, format="JPEG")
+        cropped_image_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
+        
+        return {"bounding_box": first_bbox, "cropped_image": cropped_image_base64}
+    
+    except HTTPException as e:
+        error_traceback = traceback.format_exc()
+        raise e
+    except Exception as e:
+        error_traceback = traceback.format_exc()
+        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")

models/face_classifier.py

@@ -12,7 +12,7 @@ def warning_with_traceback(message, category, filename, lineno, file=None, line=
     traceback.print_stack(file=log)
     log.write(warnings.formatwarning(message, category, filename, lineno, line))
     
-warnings.showwarning = warning_with_traceback
+# warnings.showwarning = warning_with_traceback
     
     
 class FaceSegmentationModel:

procedures.py

@@ -0,0 +1,61 @@
+# procedures.py
+
+import base64
+import io
+import traceback
+from abc import ABC, abstractmethod
+from typing import Optional, Union
+from utils import handlers
+from fastapi import File, Form, HTTPException, UploadFile, Response
+from huggingface_hub import hf_hub_download
+from ultralytics import YOLO
+from supervision import Detections
+
+async def detect_face(
+    model_path: str,    
+    file: Optional[UploadFile], 
+    base64_string: Optional[str],
+    
+    ):
+    try:
+        if file is None and base64_string is None:
+            raise HTTPException(status_code=400, detail="No input data provided")
+        base64_handler = handlers.Base64Handler()
+        image_handler = handlers.ImageFileHandler(successor=base64_handler)
+        input_data: Union[UploadFile, str] = file if file is not None else base64_string
+        print(input_data)
+        pil_image = await image_handler.handle(input_data)
+        if pil_image is None:
+            raise HTTPException(status_code=400, detail="Unsupported file type")
+        
+        # Load the YOLO model
+
+        model = YOLO(model_path)
+        
+        # Inference using the pil image
+        output = model(pil_image)
+        results = Detections.from_ultralytics(output[0])
+        
+        if len(results) == 0:
+            raise HTTPException(status_code=404, detail="No face detected")
+        
+        # Get the first bounding box
+        first_bbox = results[0].xyxy[0].tolist()
+        
+        # Crop the image using the bounding box
+        x_min, y_min, x_max, y_max = map(int, first_bbox)
+        cropped_image = pil_image.crop((x_min, y_min, x_max, y_max))
+        
+        # Convert cropped image to Base64
+        buffered = io.BytesIO()
+        cropped_image.save(buffered, format="JPEG")
+        cropped_image_base64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
+        
+        return {"bounding_box": first_bbox, "cropped_image": cropped_image_base64}
+    
+    except HTTPException as e:
+        error_traceback = traceback.format_exc()
+        raise e
+    except Exception as e:
+        error_traceback = traceback.format_exc()
+        raise HTTPException(status_code=500, detail=f"An error occurred: {str(e)}")

utils/handlers.py

@@ -0,0 +1,46 @@
+from abc import ABC, abstractmethod
+import base64
+import io
+from typing import Optional, Union
+from fastapi import UploadFile
+from PIL import Image
+
+class FileHandler(ABC):
+    def __init__(self, successor: Optional['FileHandler'] = None):
+        self._successor = successor
+
+    @abstractmethod
+    async def handle(self, input_data: Union[UploadFile, str]) -> Optional[Image.Image]:
+        if self._successor:
+            return await self._successor.handle(input_data)
+        return None
+
+class ImageFileHandler(FileHandler):
+    async def handle(self, input_data: Union[UploadFile, str]) -> Optional[Image.Image]:
+        # print(id(UploadFile))
+        # print(id(input_data.__class__))
+        # print(f"ImageFileHandler received: {type(input_data)}")
+        # print(f"Module of input_data: {input_data.__class__.__module__}")
+        # print(isinstance(input_data, UploadFile))
+        if hasattr(input_data, 'read') and hasattr(input_data, 'filename'):
+            print("Handling UploadFile")
+            try:
+                image_file = await input_data.read()
+                return Image.open(io.BytesIO(image_file)).convert("RGB")
+            except Exception as e:
+                print(f"Error processing UploadFile: {e}")
+                return None
+        return await super().handle(input_data)
+
+class Base64Handler(FileHandler):
+    async def handle(self, input_data: Union[UploadFile, str]) -> Optional[Image.Image]:
+        print(f"Base64Handler received: {type(input_data)}")
+        if isinstance(input_data, str):
+            print("Handling Base64 string")
+            try:
+                decoded_data = base64.b64decode(input_data)
+                # Handle Base64 decoded data (e.g., detect face in the decoded image)
+                return Image.open(io.BytesIO(decoded_data)).convert("RGB")
+            except Exception:
+                pass
+        return await super().handle(input_data)