65 lines
2.0 KiB
Python
65 lines
2.0 KiB
Python
_V = "/runpod-volume"
|
|
import os
|
|
if os.path.isdir(_V):
|
|
os.environ.setdefault("HF_HOME", os.path.join(_V, "huggingface"))
|
|
|
|
import io
|
|
import base64
|
|
import runpod
|
|
from PIL import Image
|
|
|
|
_model = None
|
|
|
|
|
|
def _load_model():
|
|
global _model
|
|
if _model is None:
|
|
from ultralytics import YOLO
|
|
weights = os.environ.get("YOLO_MODEL", "yolov8n.pt")
|
|
_model = YOLO(weights)
|
|
return _model
|
|
|
|
|
|
def handler(event):
|
|
try:
|
|
data = (event or {}).get("input") or {}
|
|
image_b64 = data.get("image")
|
|
if not image_b64:
|
|
return {"error": "ValueError: missing 'image' in input"}
|
|
|
|
if "," in image_b64:
|
|
image_b64 = image_b64.split(",", 1)[1]
|
|
image_bytes = base64.b64decode(image_b64)
|
|
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
|
|
|
|
# Custom construction models are often less confident than COCO — use a low
|
|
# detection floor (env YOLO_CONF, default 0.15) so real brick/pipe/steel hits
|
|
# come through instead of being dropped by ultralytics' default 0.25.
|
|
conf = float(data.get("conf") or os.environ.get("YOLO_CONF", "0.15"))
|
|
model = _load_model()
|
|
results = model.predict(image, conf=conf, verbose=False)
|
|
|
|
detections = []
|
|
for result in results:
|
|
boxes = getattr(result, "boxes", None)
|
|
if boxes is None:
|
|
continue
|
|
names = result.names if getattr(result, "names", None) else model.names
|
|
for box in boxes:
|
|
cls = int(box.cls[0])
|
|
conf = float(box.conf[0])
|
|
xyxy = box.xyxy[0].tolist()
|
|
x1, y1, x2, y2 = [float(v) for v in xyxy]
|
|
detections.append({
|
|
"name": str(names[cls]),
|
|
"confidence": conf,
|
|
"box": {"x1": x1, "y1": y1, "x2": x2, "y2": y2},
|
|
})
|
|
|
|
return {"detections": detections}
|
|
except Exception as exc:
|
|
return {"error": f"{type(exc).__name__}: {exc}"}
|
|
|
|
|
|
runpod.serverless.start({"handler": handler})
|