File size: 2,192 Bytes
bb26976 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 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)}") |