eWaSR Script for OAK

[JuanFuriaz]

Hello @tersekmatija,

Thanks for the awesome project!

I'm trying to run ewasr_resnet18.blob in my OAK-D camera. For this, I first exported the model using export.py and then created a script, which follows the general ideal of prediction.py.

Nevertheless, the outputs of the OAK-D inference differ substantially from the ones of prediction.py (ran on my PC). Can you please help me out? We can surely put later the code in the repository ;)

My code:

from pathlib import Path
import cv2
import depthai as dai
import numpy as np


# Parameters
dir = "/your/directory/to/ewasr_resnet18.blob"
fps = 20

shape_rgb = (3, 384, 512)

# Load model
nnBlobPath = str((Path(__file__).parent / Path(dir)).resolve().absolute())
if not Path(nnBlobPath).exists():
    import sys
    raise FileNotFoundError(f'Required file/s not found, please run "{sys.executable} install_requirements.py"')

pipeline = dai.Pipeline()

# Define RGB camera
camRgb = pipeline.create(dai.node.ColorCamera)
camRgb.setPreviewSize(shape_rgb[2], shape_rgb[1])
camRgb.setFps(fps)
camRgb.setResolution(dai.ColorCameraProperties.SensorResolution.THE_1080_P)
camRgb.setInterleaved(False)
camRgb.setColorOrder(dai.ColorCameraProperties.ColorOrder.BGR)
camRgb.setFp16(True)  # Model requires FP16 input

# Define neural network
nn = pipeline.create(dai.node.NeuralNetwork)
nn.setBlobPath(nnBlobPath)
nn.setNumInferenceThreads(2)
camRgb.preview.link(nn.input)

# xout for rgb and neural network
xoutNN = pipeline.create(dai.node.XLinkOut)
xoutNN.setStreamName("nn")
nn.out.link(xoutNN.input)

xoutRgb = pipeline.create(dai.node.XLinkOut)
xoutRgb.setStreamName("rgb")
nn.passthrough.link(xoutRgb.input)


def interpol_mask(frame, height, width):
    # Interpolation using CV2
    mask = cv2.resize(frame, (width, height), interpolation=cv2.INTER_LINEAR)
    return mask


with dai.Device(pipeline) as device:
    # Output queues will be used to get the rgb frames and nn data from the outputs defined above
    previewQueue = device.getOutputQueue(name="rgb", maxSize=4, blocking=False)
    networkQueue = device.getOutputQueue(name="nn", maxSize=4, blocking=False)

    while True:
        inRgb = previewQueue.get()
        in_nn = networkQueue.tryGet()

        # Model needs FP16 so we have to convert color frame back to U8 on the host
        frame = np.array(inRgb.getData()).view(np.float16).reshape(shape_rgb).transpose(1, 2, 0).astype(np.uint8).copy()
        # Prediction mask with extended image size
        pred_masks = np.zeros((shape_rgb[1], shape_rgb[2], 3))
        if in_nn is not None:
            # Get segmentation mask from nn and interpolate
            pred_origin = np.reshape(in_nn.getLayerFp16('prediction'), (3, 96, 128))
            pred_mask_1, pred_mask_2, pred_mask_3 = interpol_mask(pred_origin[0], shape_rgb[1], shape_rgb[2]), \
                                                    interpol_mask(pred_origin[1], shape_rgb[1], shape_rgb[2]), \
                                                    interpol_mask(pred_origin[2], shape_rgb[1], shape_rgb[2])
            pred_masks = np.zeros((shape_rgb[1], shape_rgb[2], 3))
            pred_masks[:, :, 0], pred_masks[:, :, 1], pred_masks[:, :, 2] = pred_mask_1, pred_mask_2, pred_mask_3

            # Argmax to get a single channel mask, then transform it to 3 channels to allow color overlay
            pred_masks_arg = np.argmax(pred_masks, axis=2)
            pred_masks[pred_masks_arg == 0] = [247, 195, 37]
            pred_masks[pred_masks_arg == 1] = [41, 167, 224]
            pred_masks[pred_masks_arg == 2] = [90, 75, 164]

        if pred_masks is not None:
            # convert to uint8 to use addWeighted
            pred_masks = pred_masks.astype(np.uint8)
            frame_mix = cv2.addWeighted(frame, 0.3, pred_masks, 0.7, 0)
            cv2.imshow("mask ", frame_mix)
        cv2.imshow("rgb", frame)

        if cv2.waitKey(1) == ord('q'):
            break

Thanks in addition!

[MoffKalast]

This is the code that worked for me, albeit on an Oak-1: #4 (comment)

[tersekmatija]

@JuanFuriaz
Hey, sorry for a late reply. Yes, please use the code from #4 .

In general, you shouldn't have to set setFp16. I'd also recommend softmaxing/argmaxing before upscaling.