procedures.py

# 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)}")