# Imports import sys import time from pathlib import Path import cv2 import numpy as np import torch from IPython.display import Markdown, display from fastseg import MobileV3Large from openvino.runtime import Core sys.path.append("../utils") from notebook_utils import CityScapesSegmentation, segmentation_map_to_image, viz_result_image # Settings IMAGE_WIDTH = 1024 # Suggested values: 2048, 1024 or 512. The minimum width is 512. # Set IMAGE_HEIGHT manually for custom input sizes. Minimum height is 512 IMAGE_HEIGHT = 1024 if IMAGE_WIDTH == 2048 else 512 DIRECTORY_NAME = "model" BASE_MODEL_NAME = DIRECTORY_NAME + f"/fastseg{IMAGE_WIDTH}" # Paths where PyTorch, ONNX and OpenVINO IR models will be stored model_path = Path(BASE_MODEL_NAME).with_suffix(".pth") onnx_path = model_path.with_suffix(".onnx") ir_path = model_path.with_suffix(".xml") # Download the Fastseg Model print("Downloading the Fastseg model (if it has not been downloaded before)....") model = MobileV3Large.from_pretrained().cpu().eval() print("Loaded PyTorch Fastseg model") # Save the model model_path.parent.mkdir(exist_ok=True) torch.save(model.state_dict(), str(model_path)) print(f"Model saved at {model_path}") # Convert PyTorch model to ONNX if not onnx_path.exists(): dummy_input = torch.randn(1, 3, IMAGE_HEIGHT, IMAGE_WIDTH) # For the Fastseg model, setting do_constant_folding to False is required # for PyTorch>1.5.1 torch.onnx.export( model, dummy_input, onnx_path, opset_version=11, do_constant_folding=False, ) print(f"ONNX model exported to {onnx_path}.") else: print(f"ONNX model {onnx_path} already exists.") # Convert ONNX Model to OpenVINO IR Format # Construct the command for Model Optimizer mo_command = f"""mo --input_model "{onnx_path}" --input_shape "[1,3, {IMAGE_HEIGHT}, {IMAGE_WIDTH}]" --mean_values="[123.675, 116.28 , 103.53]" --scale_values="[58.395, 57.12 , 57.375]" --data_type FP16 --output_dir "{model_path.parent}" """ mo_command = " ".join(mo_command.split()) print("Model Optimizer command to convert the ONNX model to OpenVINO:") display(Markdown(f"`{mo_command}`")) if not ir_path.exists(): print("Exporting ONNX model to IR... This may take a few minutes.") mo_result = %sx $mo_command print("\n".join(mo_result)) else: print(f"IR model {ir_path} already exists.") # Show results: Load and Preprocess an Input Image def normalize(image: np.ndarray) -> np.ndarray: """ Normalize the image to the given mean and standard deviation for CityScapes models. """ image = image.astype(np.float32) mean = (0.485, 0.456, 0.406) std = (0.229, 0.224, 0.225) image /= 255.0 image -= mean image /= std return image image_filename = "data/street.jpg" image = cv2.cvtColor(cv2.imread(image_filename), cv2.COLOR_BGR2RGB) resized_image = cv2.resize(image, (IMAGE_WIDTH, IMAGE_HEIGHT)) normalized_image = normalize(resized_image) # Convert the resized images to network input shape input_image = np.expand_dims(np.transpose(resized_image, (2, 0, 1)), 0) normalized_input_image = np.expand_dims(np.transpose(normalized_image, (2, 0, 1)), 0) # ONNX Model in Inference Engine # Load network to Inference Engine ie = Core() model_onnx = ie.read_model(model=onnx_path) compiled_model_onnx = ie.compile_model(model=model_onnx, device_name="CPU") output_layer_onnx = compiled_model_onnx.output(0) # Run inference on the input image res_onnx = compiled_model_onnx([normalized_input_image])[output_layer_onnx] # Convert network result to segmentation map and display the result result_mask_onnx = np.squeeze(np.argmax(res_onnx, axis=1)).astype(np.uint8) viz_result_image( image, segmentation_map_to_image(result_mask_onnx, CityScapesSegmentation.get_colormap()), resize=True, ) # IR Model in Inference Engine # Load the network in Inference Engine ie = Core() model_ir = ie.read_model(model=ir_path) compiled_model_ir = ie.compile_model(model=model_ir, device_name="CPU") # Get input and output layers output_layer_ir = compiled_model_ir.output(0) # Run inference on the input image res_ir = compiled_model_ir([input_image])[output_layer_ir] result_mask_ir = np.squeeze(np.argmax(res_ir, axis=1)).astype(np.uint8) viz_result_image( image, segmentation_map_to_image(result=result_mask_ir, colormap=CityScapesSegmentation.get_colormap()), resize=True, ) # PyTorch Comparison with torch.no_grad(): result_torch = model(torch.as_tensor(normalized_input_image).float()) result_mask_torch = torch.argmax(result_torch, dim=1).squeeze(0).numpy().astype(np.uint8) viz_result_image( image, segmentation_map_to_image(result=result_mask_torch, colormap=CityScapesSegmentation.get_colormap()), resize=True, ) # Performance Comparison num_images = 20 start = time.perf_counter() for _ in range(num_images): compiled_model_onnx([normalized_input_image]) end = time.perf_counter() time_onnx = end - start print( f"ONNX model in Inference Engine/CPU: {time_onnx/num_images:.3f} " f"seconds per image, FPS: {num_images/time_onnx:.2f}" ) start = time.perf_counter() for _ in range(num_images): compiled_model_ir([input_image]) end = time.perf_counter() time_ir = end - start print( f"IR model in Inference Engine/CPU: {time_ir/num_images:.3f} " f"seconds per image, FPS: {num_images/time_ir:.2f}" ) with torch.no_grad(): start = time.perf_counter() for _ in range(num_images): model(torch.as_tensor(input_image).float()) end = time.perf_counter() time_torch = end - start print( f"PyTorch model on CPU: {time_torch/num_images:.3f} seconds per image, " f"FPS: {num_images/time_torch:.2f}" ) if "GPU" in ie.available_devices: compiled_model_onnx_gpu = ie.compile_model(model=model_onnx, device_name="GPU") start = time.perf_counter() for _ in range(num_images): compiled_model_onnx_gpu([input_image]) end = time.perf_counter() time_onnx_gpu = end - start print( f"ONNX model in Inference Engine/GPU: {time_onnx_gpu/num_images:.3f} " f"seconds per image, FPS: {num_images/time_onnx_gpu:.2f}" ) compiled_model_ir_gpu = ie.compile_model(model=model_ir, device_name="GPU") start = time.perf_counter() for _ in range(num_images): compiled_model_ir_gpu([input_image]) end = time.perf_counter() time_ir_gpu = end - start print( f"IR model in Inference Engine/GPU: {time_ir_gpu/num_images:.3f} " f"seconds per image, FPS: {num_images/time_ir_gpu:.2f}" ) # Show Device Information devices = ie.available_devices for device in devices: device_name = ie.get_property(device_name=device, name="FULL_DEVICE_NAME") print(f"{device}: {device_name}")
Preview:
downloadDownload PNG
downloadDownload JPEG
downloadDownload SVG
Tip: You can change the style, width & colours of the snippet with the inspect tool before clicking Download!
Click to optimize width for Twitter