Fucius's picture
Upload 422 files
df6c67d verified
raw
history blame
12.3 kB
import json
import threading
import time
import traceback
from typing import Callable, Union
import cv2
import numpy as np
import supervision as sv
from PIL import Image
import inference.core.entities.requests.inference
from inference.core.active_learning.middlewares import (
NullActiveLearningMiddleware,
ThreadingActiveLearningMiddleware,
)
from inference.core.cache import cache
from inference.core.env import (
ACTIVE_LEARNING_ENABLED,
API_KEY,
API_KEY_ENV_NAMES,
CLASS_AGNOSTIC_NMS,
CONFIDENCE,
ENABLE_BYTE_TRACK,
ENFORCE_FPS,
IOU_THRESHOLD,
JSON_RESPONSE,
MAX_CANDIDATES,
MAX_DETECTIONS,
MODEL_ID,
STREAM_ID,
)
from inference.core.interfaces.base import BaseInterface
from inference.core.interfaces.camera.camera import WebcamStream
from inference.core.logger import logger
from inference.core.registries.roboflow import get_model_type
from inference.models.utils import get_roboflow_model
class Stream(BaseInterface):
"""Roboflow defined stream interface for a general-purpose inference server.
Attributes:
model_manager (ModelManager): The manager that handles model inference tasks.
model_registry (RoboflowModelRegistry): The registry to fetch model instances.
api_key (str): The API key for accessing models.
class_agnostic_nms (bool): Flag for class-agnostic non-maximum suppression.
confidence (float): Confidence threshold for inference.
iou_threshold (float): The intersection-over-union threshold for detection.
json_response (bool): Flag to toggle JSON response format.
max_candidates (float): The maximum number of candidates for detection.
max_detections (float): The maximum number of detections.
model (str|Callable): The model to be used.
stream_id (str): The ID of the stream to be used.
use_bytetrack (bool): Flag to use bytetrack,
Methods:
init_infer: Initialize the inference with a test frame.
preprocess_thread: Preprocess incoming frames for inference.
inference_request_thread: Manage the inference requests.
run_thread: Run the preprocessing and inference threads.
"""
def __init__(
self,
api_key: str = API_KEY,
class_agnostic_nms: bool = CLASS_AGNOSTIC_NMS,
confidence: float = CONFIDENCE,
enforce_fps: bool = ENFORCE_FPS,
iou_threshold: float = IOU_THRESHOLD,
max_candidates: float = MAX_CANDIDATES,
max_detections: float = MAX_DETECTIONS,
model: Union[str, Callable] = MODEL_ID,
source: Union[int, str] = STREAM_ID,
use_bytetrack: bool = ENABLE_BYTE_TRACK,
use_main_thread: bool = False,
output_channel_order: str = "RGB",
on_prediction: Callable = None,
on_start: Callable = None,
on_stop: Callable = None,
):
"""Initialize the stream with the given parameters.
Prints the server settings and initializes the inference with a test frame.
"""
logger.info("Initializing server")
self.frame_count = 0
self.byte_tracker = sv.ByteTrack() if use_bytetrack else None
self.use_bytetrack = use_bytetrack
if source == "webcam":
stream_id = 0
else:
stream_id = source
self.stream_id = stream_id
if self.stream_id is None:
raise ValueError("STREAM_ID is not defined")
self.model_id = model
if not self.model_id:
raise ValueError("MODEL_ID is not defined")
self.api_key = api_key
self.active_learning_middleware = NullActiveLearningMiddleware()
if isinstance(model, str):
self.model = get_roboflow_model(model, self.api_key)
if ACTIVE_LEARNING_ENABLED:
self.active_learning_middleware = (
ThreadingActiveLearningMiddleware.init(
api_key=self.api_key,
model_id=self.model_id,
cache=cache,
)
)
self.task_type = get_model_type(
model_id=self.model_id, api_key=self.api_key
)[0]
else:
self.model = model
self.task_type = "unknown"
self.class_agnostic_nms = class_agnostic_nms
self.confidence = confidence
self.iou_threshold = iou_threshold
self.max_candidates = max_candidates
self.max_detections = max_detections
self.use_main_thread = use_main_thread
self.output_channel_order = output_channel_order
self.inference_request_type = (
inference.core.entities.requests.inference.ObjectDetectionInferenceRequest
)
self.webcam_stream = WebcamStream(
stream_id=self.stream_id, enforce_fps=enforce_fps
)
logger.info(
f"Streaming from device with resolution: {self.webcam_stream.width} x {self.webcam_stream.height}"
)
self.on_start_callbacks = []
self.on_stop_callbacks = [
lambda: self.active_learning_middleware.stop_registration_thread()
]
self.on_prediction_callbacks = []
if on_prediction:
self.on_prediction_callbacks.append(on_prediction)
if on_start:
self.on_start_callbacks.append(on_start)
if on_stop:
self.on_stop_callbacks.append(on_stop)
self.init_infer()
self.preproc_result = None
self.inference_request_obj = None
self.queue_control = False
self.inference_response = None
self.stop = False
self.frame = None
self.frame_cv = None
self.frame_id = None
logger.info("Server initialized with settings:")
logger.info(f"Stream ID: {self.stream_id}")
logger.info(f"Model ID: {self.model_id}")
logger.info(f"Enforce FPS: {enforce_fps}")
logger.info(f"Confidence: {self.confidence}")
logger.info(f"Class Agnostic NMS: {self.class_agnostic_nms}")
logger.info(f"IOU Threshold: {self.iou_threshold}")
logger.info(f"Max Candidates: {self.max_candidates}")
logger.info(f"Max Detections: {self.max_detections}")
self.run_thread()
def on_start(self, callback):
self.on_start_callbacks.append(callback)
unsubscribe = lambda: self.on_start_callbacks.remove(callback)
return unsubscribe
def on_stop(self, callback):
self.on_stop_callbacks.append(callback)
unsubscribe = lambda: self.on_stop_callbacks.remove(callback)
return unsubscribe
def on_prediction(self, callback):
self.on_prediction_callbacks.append(callback)
unsubscribe = lambda: self.on_prediction_callbacks.remove(callback)
return unsubscribe
def init_infer(self):
"""Initialize the inference with a test frame.
Creates a test frame and runs it through the entire inference process to ensure everything is working.
"""
frame = Image.new("RGB", (640, 640), color="black")
self.model.infer(
frame, confidence=self.confidence, iou_threshold=self.iou_threshold
)
self.active_learning_middleware.start_registration_thread()
def preprocess_thread(self):
"""Preprocess incoming frames for inference.
Reads frames from the webcam stream, converts them into the proper format, and preprocesses them for
inference.
"""
webcam_stream = self.webcam_stream
webcam_stream.start()
# processing frames in input stream
try:
while True:
if webcam_stream.stopped is True or self.stop:
break
else:
self.frame_cv, frame_id = webcam_stream.read_opencv()
if frame_id > 0 and frame_id != self.frame_id:
self.frame_id = frame_id
self.frame = cv2.cvtColor(self.frame_cv, cv2.COLOR_BGR2RGB)
self.preproc_result = self.model.preprocess(self.frame_cv)
self.img_in, self.img_dims = self.preproc_result
self.queue_control = True
except Exception as e:
traceback.print_exc()
logger.error(e)
def inference_request_thread(self):
"""Manage the inference requests.
Processes preprocessed frames for inference, post-processes the predictions, and sends the results
to registered callbacks.
"""
last_print = time.perf_counter()
print_ind = 0
while True:
if self.webcam_stream.stopped is True or self.stop:
while len(self.on_stop_callbacks) > 0:
# run each onStop callback only once from this thread
cb = self.on_stop_callbacks.pop()
cb()
break
if self.queue_control:
while len(self.on_start_callbacks) > 0:
# run each onStart callback only once from this thread
cb = self.on_start_callbacks.pop()
cb()
self.queue_control = False
frame_id = self.frame_id
inference_input = np.copy(self.frame_cv)
start = time.perf_counter()
predictions = self.model.predict(
self.img_in,
)
predictions = self.model.postprocess(
predictions,
self.img_dims,
class_agnostic_nms=self.class_agnostic_nms,
confidence=self.confidence,
iou_threshold=self.iou_threshold,
max_candidates=self.max_candidates,
max_detections=self.max_detections,
)[0]
self.active_learning_middleware.register(
inference_input=inference_input,
prediction=predictions.dict(by_alias=True, exclude_none=True),
prediction_type=self.task_type,
)
if self.use_bytetrack:
detections = sv.Detections.from_roboflow(
predictions.dict(by_alias=True, exclude_none=True)
)
detections = self.byte_tracker.update_with_detections(detections)
if detections.tracker_id is None:
detections.tracker_id = np.array([], dtype=int)
for pred, detect in zip(predictions.predictions, detections):
pred.tracker_id = int(detect[4])
predictions.frame_id = frame_id
predictions = predictions.dict(by_alias=True, exclude_none=True)
self.inference_response = predictions
self.frame_count += 1
for cb in self.on_prediction_callbacks:
if self.output_channel_order == "BGR":
cb(predictions, self.frame_cv)
else:
cb(predictions, np.asarray(self.frame))
current = time.perf_counter()
self.webcam_stream.max_fps = 1 / (current - start)
logger.debug(f"FPS: {self.webcam_stream.max_fps:.2f}")
if time.perf_counter() - last_print > 1:
print_ind = (print_ind + 1) % 4
last_print = time.perf_counter()
def run_thread(self):
"""Run the preprocessing and inference threads.
Starts the preprocessing and inference threads, and handles graceful shutdown on KeyboardInterrupt.
"""
preprocess_thread = threading.Thread(target=self.preprocess_thread)
preprocess_thread.start()
if self.use_main_thread:
self.inference_request_thread()
else:
# start a thread that looks for the predictions
# and call the callbacks
inference_request_thread = threading.Thread(
target=self.inference_request_thread
)
inference_request_thread.start()