Python Linux/Jetson Code Example

You can use the provided Python example code which includes image reading, model loading, model prediction, and output drawing.

After downloading, extract the folder. If you are using a Docker Image, ensure the files are placed under /home/appuser/workdir, which is the local directory when the container starts.

Then, run the following command to execute the Python script and make predictions using the provided example model and image. The results will be output in the same folder as output_with_masks_classes_scores.png.

python3 example.py

You can change the file paths in example.py to use different images and deep learning models.

MODEL_ZIP = 'kp1.zip'
IMG_PATH = 'kp.png'

You can also start with a new Python file. First, you need to import the daoai_vision library, which is our DaoAI World Python SDK.

import daoai_vision as dv

The following libraries may also be helpful:

import cv2
import numpy as np
import matplotlib.pyplot as plt

Loading a Model

MODEL_ZIP = 'kp_fast_model.zip'  # DaoAI world model file path
DEVICE = 'gpu'  # Hardware option: cpu or gpu. Using cpu will significantly slow down the model prediction.
model = dv.get_model(model_zip=MODEL_ZIP, device=DEVICE)

Loading an Image and Running Model Prediction

IMG_PATH = 'test_image.png'  # Path to the image being read
results = model.infer(IMG_PATH)  # Call the model prediction function

Prediction Results

boxes     = results.boxes     # [N][x1, y1, x2, y2]: bounding box formed by two points
classes   = results.classes   # [N] Model class index values
labels    = results.labels    # [N] Model class strings; you can use labels[classes[i]] to access the corresponding class string
scores    = results.scores    # [N] Confidence score of the prediction
masks     = results.masks     # [N][H][W] Model mask pixels
keypoints = results.keypoints # [N][x, y, score] Model keypoints
sem_masks = results.sem_seg   # Pixel masks

You can use these results as needed.

Drawing Bounding Boxes

# Draw bounding boxes
if(boxes is not None):
    for box in boxes:
        cv2.rectangle(img, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (0, 255, 0), 2)

Drawing Keypoints

if keypoints is not None:
    for obj in keypoints:
        for kp in obj:
            if float(kp[2]) > 0.5:
                cv2.circle(img, (int(kp[0]), int(kp[1])), 2, (0, 0, 255), 2)

Drawing Masks, Classes, and Scores

# Draw masks, classes, and scores
for i in range(len(boxes)):
    if masks is not None:
        mask = masks[i]

        # Resize the mask to image size
        mask_resized = cv2.resize(mask.astype(np.uint8), (img.shape[1], img.shape[0]))
        mask_rgb = np.zeros_like(img)
        mask_rgb[mask_resized > 0.5] = (255, 0, 255)  # Set the mask color
        # Draw mask
        img = cv2.addWeighted(img, 1, mask_rgb, 0.5, 0)

    cls = classes[i]
    score = scores[i]
    label = labels[cls]

    # Get current bounding box coordinates
    bbox = boxes[i]
    x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])

    # Draw class and score
    text = f'{cls} {label}: {score:.2f}'
    cv2.putText(img, text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

Using the Example Code Prediction Results: