def __init__(self, *, cfg, dataset_type: callable, data_root: Path, sub_datasets: Tuple, split: Split = Split.Training): """ :param data_root: :type data_root: :param sub_datasets: :type sub_datasets: :param transform: :type transform: :param target_transform: :type target_transform: :param split: :type split: :return: :rtype: """ assert len(sub_datasets) > 0, "No data found!" img_transform = SSDTransform( image_size=cfg.input.image_size, pixel_mean=cfg.input.pixel_mean, split=split, ) if split == Split.Training: annotation_transform = SSDAnnotationTransform( image_size=cfg.input.image_size, priors_cfg=cfg.model.box_head.priors, center_variance=cfg.model.box_head.center_variance, size_variance=cfg.model.box_head.size_variance, iou_threshold=cfg.model.box_head.iou_threshold, ) else: annotation_transform = None datasets = [] for dataset_name in sub_datasets: datasets.append( dataset_type( data_root=data_root, dataset_name=dataset_name, split=split, img_transform=img_transform, annotation_transform=annotation_transform, )) # for testing, return a list of datasets if not split == Split.Training: self.sub_datasets = datasets else: dataset = datasets[0] if len(datasets) > 1: dataset = ConcatDataset(datasets) self.sub_datasets = [dataset]
def run_webcam_demo( cfg: NOD, input_cfg: NOD, categories: List, model_ckpt: Path, score_threshold: float = 0.7, window_name: str = "SSD", ): """ :param categories: :type categories: :param cfg: :type cfg: :param model_ckpt: :type model_ckpt: :param score_threshold: :type score_threshold: :param window_name: :type window_name: :return: :rtype: """ cpu_device = torch.device("cpu") transforms = SSDTransform(input_cfg.image_size, input_cfg.pixel_mean, split=Split.Testing) model = SingleShotDectectionNms(cfg) checkpointer = CheckPointer(model, save_dir=ensure_existence( PROJECT_APP_PATH.user_data / "results")) checkpointer.load(model_ckpt, use_latest=model_ckpt is None) print( f"Loaded weights from {model_ckpt if model_ckpt else checkpointer.get_checkpoint_file()}" ) model.post_init() model.to(global_torch_device()) with TorchEvalSession(model): for image in tqdm(frame_generator(cv2.VideoCapture(0))): result = model( transforms(image)[0].unsqueeze(0).to(global_torch_device())) height, width, *_ = image.shape result.boxes[:, 0::2] *= width / result.img_width.cpu().item() result.boxes[:, 1::2] *= height / result.img_height.cpu().item() (boxes, labels, scores) = ( result.boxes.to(cpu_device).numpy(), result.labels.to(cpu_device).numpy(), result.scores.to(cpu_device).numpy(), ) indices = scores > score_threshold cv2.namedWindow(window_name, cv2.WINDOW_NORMAL) cv2.imshow( window_name, draw_bounding_boxes( image, boxes[indices], labels=labels[indices], scores=scores[indices], categories=categories, score_font=ImageFont.truetype( PACKAGE_DATA_PATH / "Lato-Regular.ttf", 24, ), ).astype(numpy.uint8), ) if cv2.waitKey(1) == 27: break # esc to quit
def export_detection_model( cfg: NOD, model_ckpt: Path, model_export_path: Path = Path("torch_model"), verbose: bool = True, onnx_export: bool = False, strict_jit: bool = False, ) -> None: """ :param verbose: :type verbose: :param cfg: :type cfg: :param model_ckpt: :type model_ckpt: :param model_export_path: :type model_export_path: :return: :rtype: """ model = SingleShotDectectionNms(cfg) checkpointer = CheckPointer( model, save_dir=ensure_existence(PROJECT_APP_PATH.user_data / "results") ) checkpointer.load(model_ckpt, use_latest=model_ckpt is None) print( f"Loaded weights from {model_ckpt if model_ckpt else checkpointer.get_checkpoint_file()}" ) model.post_init() model.to(global_torch_device()) transforms = SSDTransform( cfg.input.image_size, cfg.input.pixel_mean, split=Split.Testing ) model.eval() pre_quantize_model = False if pre_quantize_model: # Accuracy may drop! if True: model = quantization.quantize_dynamic(model, dtype=torch.qint8) else: pass # model = quantization.quantize(model) frame_g = frame_generator(cv2.VideoCapture(0)) for image in tqdm(frame_g): example_input = (transforms(image)[0].unsqueeze(0).to(global_torch_device()),) try: traced_script_module = torch.jit.script( model, # example_input, ) exp_path = model_export_path.with_suffix(".compiled") traced_script_module.save(str(exp_path)) print(f"Traced Ops used {torch.jit.export_opnames(traced_script_module)}") sprint( f"Successfully exported JIT Traced model at {exp_path}", color="green" ) except Exception as e_i: sprint(f"Torch JIT Trace export does not work!, {e_i}", color="red") break
def export_detection_model( cfg: NOD, model_checkpoint: Path, model_export_path: Path = Path("torch_model"), verbose: bool = True, onnx_export: bool = False, strict_jit: bool = False, ) -> None: """ :param verbose: :type verbose: :param cfg: :type cfg: :param model_checkpoint: :type model_checkpoint: :param model_export_path: :type model_export_path: :return: :rtype:""" model = SingleShotDetection(cfg) checkpointer = CheckPointer(model, save_dir=ensure_existence( PROJECT_APP_PATH.user_data / "results")) checkpointer.load(model_checkpoint, use_latest=model_checkpoint is None) print( f"Loaded weights from {model_checkpoint if model_checkpoint else checkpointer.get_checkpoint_file()}" ) model.post_init() model.to(global_torch_device()) transforms = SSDTransform(cfg.input.image_size, cfg.input.pixel_mean, split=SplitEnum.testing) model.eval() # Important! fuse_quantize_model = False if fuse_quantize_model: modules_to_fuse = [ ["conv", "bn", "relu"] ] # Names of modules to fuse, maybe supply directly for architecture class/declaration model = torch.quantization.fuse_modules( model, modules_to_fuse=modules_to_fuse, inplace=False) pre_quantize_model = False if pre_quantize_model: # Accuracy may drop! if True: model = quantization.quantize_dynamic(model, dtype=torch.qint8) else: pass # model = quantization.quantize(model) frame_g = frame_generator(cv2.VideoCapture(0)) for image in tqdm(frame_g): example_input = (transforms(image)[0].unsqueeze(0).to( global_torch_device()), ) try: if onnx_export: exp_path = model_export_path.with_suffix(".onnx") output = onnx.export( model, example_input, str(exp_path), verbose=verbose, # export_params=True, # store the trained parameter weights inside the model file # opset_version=10, # the onnx version to export the model to # do_constant_folding=True, # wether to execute constant folding for optimization # input_names=["input"], # the model's input names # output_names=["output"], # the model's output names # dynamic_axes={ # "input": {0: "batch_size"}, # variable lenght axes # "output": {0: "batch_size"}, # } ) sprint(f"Successfully exported ONNX model at {exp_path}", color="blue") else: raise Exception("Just trace instead, ignore exception") except Exception as e: sprint(f"Torch ONNX export does not work, {e}", color="red") try: traced_script_module = torch.jit.trace( model, example_input, # strict=strict_jit, check_inputs=( transforms(next(frame_g))[0].unsqueeze(0).to( global_torch_device()), transforms(next(frame_g))[0].unsqueeze(0).to( global_torch_device()), ), ) exp_path = model_export_path.with_suffix(".traced") traced_script_module.save(str(exp_path)) print( f"Traced Ops used {torch.jit.export_opnames(traced_script_module)}" ) sprint( f"Successfully exported JIT Traced model at {exp_path}", color="green", ) except Exception as e_i: sprint(f"Torch JIT Trace export does not work!, {e_i}", color="red") break """
def run_webcam_demo( cfg: NOD, categories: Sequence[str], model_checkpoint: Path, score_threshold: float = 0.5, window_name: str = "SSD", ): """ :param categories: :type categories: :param cfg: :type cfg: :param model_checkpoint: :type model_checkpoint: :param score_threshold: :type score_threshold: :param window_name: :type window_name: :return: :rtype:""" cpu_device = torch.device("cpu") transforms = SSDTransform(cfg.input.image_size, cfg.input.pixel_mean, split=SplitEnum.testing) model = SingleShotDetection(cfg) checkpointer = CheckPointer(model, save_dir=ensure_existence( PROJECT_APP_PATH.user_data / "results")) checkpointer.load(model_checkpoint, use_latest=model_checkpoint is None) print( f"Loaded weights from {model_checkpoint if model_checkpoint else checkpointer.get_checkpoint_file()}" ) model.post_init() model.to(global_torch_device()) with TorchEvalSession(model): for infos in tqdm(DictUnityEnvironment(connect_to_running=True)): info = next(iter(infos.values())) new_images = extract_all_cameras(info) image = next(iter(new_images.values()))[..., :3][..., ::-1] image = gamma_correct_float_to_byte(image) result = model( transforms(image)[0].unsqueeze(0).to(global_torch_device()))[0] height, width, *_ = image.shape result["boxes"][:, 0::2] *= width / result["img_width"] result["boxes"][:, 1::2] *= height / result["img_height"] (boxes, labels, scores) = ( result["boxes"].to(cpu_device).numpy(), result["labels"].to(cpu_device).numpy(), result["scores"].to(cpu_device).numpy(), ) indices = scores > score_threshold if show_image( draw_bounding_boxes( image, boxes[indices], labels=labels[indices], scores=scores[indices], categories=categories, ).astype(numpy.uint8), window_name, wait=1, ): break # esc to quit
def run_traced_webcam_demo( input_cfg: NOD, categories: List, score_threshold: float = 0.7, window_name: str = "SSD", onnx_exported: bool = False, ): """ :param onnx_exported: :type onnx_exported: :param input_cfg: :type input_cfg: :param categories: :type categories: :param score_threshold: :type score_threshold: :param window_name: :type window_name: :return: :rtype: """ pass import torch cpu_device = torch.device("cpu") transforms = SSDTransform(input_cfg.image_size, input_cfg.pixel_mean, split=Split.Testing) model = None if onnx_exported: import onnx onnx_model = onnx.load("torch_model.onnx") onnx.checker.check_model(onnx_model) import onnxruntime ort_session = onnxruntime.InferenceSession("torch_model.onnx") def to_numpy(tensor): return (tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()) x = None # compute onnxruntime output prediction ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)} ort_outs = ort_session.run(None, ort_inputs) else: import torch import io torch.jit.load("torch_model.traced") with open("torch_model.traced", "rb") as f: # Load ScriptModule from io.BytesIO object buffer = io.BytesIO(f.read()) model = torch.jit.load( buffer) # Load all tensors to the original device """ buffer.seek(0) torch.jit.load(buffer, map_location=torch.device('cpu')) # Load all tensors onto CPU, using a device buffer.seek(0) model = torch.jit.load(buffer, map_location='cpu') # Load all tensors onto CPU, using a string # Load with extra files. extra_files = torch._C.ExtraFilesMap() extra_files['foo.txt'] = 'bar' torch.jit.load('torch_model.traced', _extra_files=extra_files) print(extra_files['foo.txt']) #exit(0) """ with TorchDeviceSession( device=global_torch_device(cuda_if_available=False), model=model): with TorchEvalSession(model): for image in tqdm(frame_generator(cv2.VideoCapture(0))): result = SSDOut(*model( transforms(image)[0].unsqueeze(0).to( global_torch_device()))) height, width, *_ = image.shape result.boxes[:, 0::2] *= width / result.img_width.cpu().item() result.boxes[:, 1::2] *= height / result.img_height.cpu().item() (boxes, labels, scores) = ( result.boxes.to(cpu_device).numpy(), result.labels.to(cpu_device).numpy(), result.scores.to(cpu_device).numpy(), ) indices = scores > score_threshold cv2.namedWindow(window_name, cv2.WINDOW_NORMAL) cv2.imshow( window_name, draw_bounding_boxes( image, boxes[indices], labels=labels[indices], scores=scores[indices], categories=categories, score_font=ImageFont.truetype( PACKAGE_DATA_PATH / "Lato-Regular.ttf", 24, ), ).astype(numpy.uint8), ) if cv2.waitKey(1) == 27: break # esc to quit
def run_demo(cfg, class_names, model_ckpt, score_threshold, images_dir, output_dir): model = SingleShotDectection(cfg) checkpointer = CheckPointer(model, save_dir=ensure_existence( PROJECT_APP_PATH.user_data / "results")) checkpointer.load(model_ckpt, use_latest=model_ckpt is None) print( f"Loaded weights from {model_ckpt if model_ckpt else checkpointer.get_checkpoint_file()}" ) model.post_init() model.to(global_torch_device()) image_paths = list(images_dir.iterdir()) cpu_device = torch.device("cpu") transforms = SSDTransform(cfg.input.image_size, cfg.input.pixel_mean, split=Split.Testing) model.eval() for i, image_path in enumerate(image_paths): start = time.time() image_name = os.path.basename(image_path) image = numpy.array(Image.open(image_path).convert("RGB")) height, width = image.shape[:2] images = transforms(image)[0].unsqueeze(0) load_time = time.time() - start start = time.time() result = model(images.to(global_torch_device()))[0] inference_time = time.time() - start result.boxes[:, 0::2] *= width / result.img_width result.boxes[:, 1::2] *= height / result.img_height (boxes, labels, scores) = ( result.boxes.to(cpu_device).numpy(), result.labels.to(cpu_device).numpy(), result.scores.to(cpu_device).numpy(), ) indices = scores > score_threshold boxes, labels, scores = boxes[indices], labels[indices], scores[ indices] meters = " | ".join([ f"objects {len(boxes):02d}", f"load {round(load_time * 1000):03d}ms", f"inference {round(inference_time * 1000):03d}ms", f"FPS {round(1.0 / inference_time)}", ]) print(f"({i + 1:04d}/{len(image_paths):04d}) {image_name}: {meters}") drawn_image = draw_bounding_boxes( image, boxes, labels, scores, class_names, score_font=ImageFont.truetype( PACKAGE_DATA_PATH / "Lato-Regular.ttf", 24, ), ).astype(numpy.uint8) Image.fromarray(drawn_image).save(os.path.join(output_dir, image_name))