diff --git a/detector/docker.sh b/detector/docker.sh index 0880781..3a77cfe 100755 --- a/detector/docker.sh +++ b/detector/docker.sh @@ -28,8 +28,8 @@ fi # ========================== BUILD CONFIGURATION / IMAGE SELECTION ======================= -SEMANTIC_VERSION=0.1.3 -NODE_LIB_VERSION=0.10.12 +SEMANTIC_VERSION=0.1.4 +NODE_LIB_VERSION=0.10.13 build_args=" --build-arg NODE_LIB_VERSION=$NODE_LIB_VERSION" if [ -f /etc/nv_tegra_release ] # Check if we are on a Jetson device @@ -75,6 +75,7 @@ run_args="-it" run_args+=" -v $HOME/node_data/$DETECTOR_NAME:/data" run_args+=" -h ${HOSTNAME}_DEV" run_args+=" -e HOST=$LOOP_HOST -e ORGANIZATION=$LOOP_ORGANIZATION -e PROJECT=$LOOP_PROJECT" +run_args+=" -e USE_BACKDOOR_CONTROLS=$USE_BACKDOOR_CONTROLS" run_args+=" --name $DETECTOR_NAME" run_args+=" --runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=all" run_args+=" -p 8004:80" diff --git a/detector/yolov5_detector.py b/detector/yolov5_detector.py index 4acff1c..1038388 100644 --- a/detector/yolov5_detector.py +++ b/detector/yolov5_detector.py @@ -4,7 +4,7 @@ import re import subprocess import time -from typing import List, Optional, Tuple +from typing import Optional, Tuple import cv2 import numpy as np @@ -22,6 +22,8 @@ def __init__(self) -> None: self.yolov5: Optional[yolov5.YoLov5TRT] = None self.weight_type = os.getenv('WEIGHT_TYPE', 'FP16') assert self.weight_type in ['FP16', 'FP32', 'INT8'], 'WEIGHT_TYPE must be one of FP16, FP32, INT8' + self.log = logging.getLogger('Yolov5Detector') + self.log.setLevel(logging.INFO) def init(self) -> None: resolution = self.model_info.resolution @@ -61,7 +63,7 @@ def clip_point(x: float, y: float, img_width: int, img_height: int) -> Tuple[flo y = min(max(0, y), img_height) return x, y - def evaluate(self, image: List[np.uint8]) -> Detections: + def evaluate(self, image: np.ndarray) -> Detections: assert self.yolov5 is not None, 'init() must be executed first. Maybe loading the engine failed?!' detections = Detections() try: @@ -70,7 +72,7 @@ def evaluate(self, image: List[np.uint8]) -> Detections: im_height, im_width, _ = cv_image.shape results, inference_ms = self.yolov5.infer(cv_image) skipped_detections = [] - logging.info(f'took {inference_ms} s, overall {time.time() -t} s') + self.log.debug('took %f s, overall %f s', inference_ms, time.time() - t) for detection in results: x, y, w, h, category_idx, probability = detection category = self.model_info.categories[category_idx] @@ -80,20 +82,30 @@ def evaluate(self, image: List[np.uint8]) -> Detections: if category.type == CategoryType.Box: x, y, w, h = self.clip_box(x, y, w, h, im_width, im_height) detections.box_detections.append( - BoxDetection(category_name=category.name, x=x, y=y, width=w, height=h, category_id=category.id, - model_name=self.model_info.version, confidence=probability)) + BoxDetection(category_name=category.name, + x=round(x), + y=round(y), + width=round(x+w)-round(x), + height=round(y+h)-round(y), + category_id=category.id, + model_name=self.model_info.version, + confidence=probability)) elif category.type == CategoryType.Point: - cx, cy = (np.average([x, x + w]), np.average([y, y + h])) + cx, cy = x + w/2, y + h/2 cx, cy = self.clip_point(cx, cy, im_width, im_height) detections.point_detections.append( - PointDetection(category_name=category.name, x=int(cx), y=int(cy), category_id=category.id, - model_name=self.model_info.version, confidence=probability)) + PointDetection(category_name=category.name, + x=cx, + y=cy, + category_id=category.id, + model_name=self.model_info.version, + confidence=probability)) if skipped_detections: log_msg = '\n'.join([str(d) for d in skipped_detections]) logging.warning( f'Removed {len(skipped_detections)} small detections from result: \n{log_msg}') except Exception: - logging.exception('inference failed') + self.log.exception('inference failed') return detections def _create_engine(self, resolution: int, cat_count: int, wts_file: str) -> str: @@ -110,13 +122,13 @@ def _create_engine(self, resolution: int, cat_count: int, wts_file: str) -> str: with open('../src/config.h', 'r+') as f: content = f.read() if self.weight_type == 'INT8': - logging.info('using INT8') + self.log.info('using INT8') content = content.replace('#define USE_FP16', '#define USE_INT8') elif self.weight_type == 'FP32': - logging.info('using FP32') + self.log.info('using FP32') content = content.replace('#define USE_FP16', '#define USE_FP32') else: - logging.info('using FP16') + self.log.info('using FP16') content = re.sub('(kNumClass =) \d*', r'\1 ' + str(cat_count), content) @@ -128,7 +140,7 @@ def _create_engine(self, resolution: int, cat_count: int, wts_file: str) -> str: subprocess.run('make -j6 -Wno-deprecated-declarations', shell=True, check=True) - logging.warning('currently we assume a Yolov5 s6 model;\ + self.log.warning('currently we assume a Yolov5 s6 model;\ parameterization of the variant (s, s6, m, m6, ...) still needs to be done') # TODO parameterize variant "s6" subprocess.run( diff --git a/trainer/docker.sh b/trainer/docker.sh index dabbbe2..27be1ae 100755 --- a/trainer/docker.sh +++ b/trainer/docker.sh @@ -44,8 +44,8 @@ fi # NODE_LIB_VERSION should only be used, to build the corresponding version and deploy to docker # make sure the remote repository always has the 'latest' tag (otherwise the CI tests will fail) -SEMANTIC_VERSION=0.1.3 -NODE_LIB_VERSION=0.10.12 +SEMANTIC_VERSION=0.1.4 +NODE_LIB_VERSION=0.10.13 #image="zauberzeug/yolov5-trainer:$SEMANTIC_VERSION-nlv$NODE_LIB_VERSION" image="zauberzeug/yolov5-trainer:latest" @@ -74,6 +74,7 @@ run_args+=" -h ${HOSTNAME}_DEV" run_args+=" -e HOST=$HOST -e USERNAME=$USERNAME -e PASSWORD=$PASSWORD -e LOOP_SSL_CERT_PATH=$LOOP_SSL_CERT_PATH" run_args+=" -e BATCH_SIZE=$BATCH_SIZE -e UVICORN_RELOAD=$UVICORN_RELOAD -e RESET_POINTS=$RESET_POINTS -e KEEP_OLD_TRAININGS=$KEEP_OLD_TRAININGS" run_args+=" -e NODE_TYPE=trainer -e YOLOV5_MODE=$YOLOV5_MODE -e RESTART_AFTER_TRAINING=$RESTART_AFTER_TRAINING -e TRAINER_IDLE_TIMEOUT_SEC=$TRAINER_IDLE_TIMEOUT_SEC" +run_args+=" -e USE_BACKDOOR_CONTROLS=$USE_BACKDOOR_CONTROLS" run_args+=" --name $TRAINER_NAME" run_args+=" --runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=all --gpus all" run_args+=" --ipc host"