| """Custom inference handler for RF-DETR Threat Detection on HuggingFace Inference Endpoints.""" |
|
|
| from typing import Any, Dict, List |
| import io |
| import numpy as np |
| from PIL import Image |
|
|
|
|
| class EndpointHandler: |
| def __init__(self, path: str = ""): |
| from rfdetr import RFDETRNano |
| import os |
|
|
| weights = os.path.join(path, "checkpoint_best_total.pth") |
| self.model = RFDETRNano(resolution=960, pretrain_weights=weights) |
| self.model.optimize_for_inference() |
|
|
| self.class_map = {1: "Gun", 2: "Explosive", 3: "Grenade", 4: "Knife"} |
|
|
| def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: |
| inputs = data.get("inputs") |
| if isinstance(inputs, bytes): |
| image = Image.open(io.BytesIO(inputs)).convert("RGB") |
| elif isinstance(inputs, str): |
| import base64 |
| image = Image.open(io.BytesIO(base64.b64decode(inputs))).convert("RGB") |
| else: |
| image = inputs |
|
|
| threshold = data.get("parameters", {}).get("threshold", 0.25) |
| detections = self.model.predict(image, threshold=threshold) |
|
|
| results = [] |
| if hasattr(detections, "class_id") and len(detections.class_id) > 0: |
| for idx in range(len(detections.class_id)): |
| cid = int(detections.class_id[idx]) |
| conf = float(detections.confidence[idx]) |
| bbox = detections.xyxy[idx].tolist() |
| results.append({ |
| "label": self.class_map.get(cid, f"threat_{cid}"), |
| "score": round(conf, 4), |
| "box": { |
| "xmin": int(bbox[0]), |
| "ymin": int(bbox[1]), |
| "xmax": int(bbox[2]), |
| "ymax": int(bbox[3]), |
| }, |
| }) |
|
|
| return results |
|
|
