🛠️ [WIP] Visualization POC

[ashwinvaidya17]

What does this PR do?

import numpy as np
from PIL import Image

from model_api.models.result_types import AnomalyResult, ClassificationResult
from model_api.visualizer import Visualizer


def main() -> None:
    visualizer = Visualizer()
    image = Image.open("./data/coco128/images/train2017/000000000074.jpg")
    image = image.resize((255, 255))  # Ideally this should be within visualization

    heatmap = np.zeros((128, 128), dtype=np.uint8)
    heatmap[64:96, 64:96] = 255

    mask = np.zeros((128, 128), dtype=np.uint8)
    mask[32:96, 32:96] = 255
    mask[40:80, 0:128] = 255

    anomaly_result = AnomalyResult(
        anomaly_map=heatmap,
        pred_boxes=np.array([[0, 0, 128, 128]]),
        pred_label="Anomaly",
        pred_mask=mask,
        pred_score=0.85,
    )
    visualizer.save(image, anomaly_result, "anomaly_result.jpg")
    visualizer.save(
        image, anomaly_result, "anomaly_result_full.jpg", visualization_type="full"
    )

classification_result = ClassificationResult(
        top_labels=[(0, "label 1", 0.85), (1, "label 2", 0.15)],
        saliency_map=heatmap,
        feature_vector=np.random.rand(128),
        raw_scores=np.random.rand(128),
    )
    visualizer.save(image, classification_result, "classification_result.jpg")
    visualizer.save(
        image,
        classification_result,
        "classification_result_full.jpg",
        visualization_type="full",
    )

if __name__ == "__main__":
    main()

Anomaly

Simple

Full

Classification

Simple

Full

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

[eugene123tw]

Thanks, Ashwin! I’ve left a few questions. Appreciate it! 😃

[eugene123tw]

This seems like an interesting use case for classification. However, for detection, will it display all bounding boxes associated with a particular class label?

[sovrasov]

Agree, seems like we have generalization issues here, but those can be handled after additional review of the usecases different than anomaly

[sovrasov]

Agree, seems like we have generalization issues here, but those can be handled after additional review of the usecases different than anomaly

[sovrasov]

It's a question if we have to use pillow here: wrappers expect a numpy array, we have to agree on image representation. numpy + opencv provide a better alignment with c++, while pillow is more convenient from the python API perspective

[eugene123tw]

@ashwinvaidya17 , in the current Visualizer._generate_simple design, the task type is not considered, which leads to inconsistent visualizations. For instance, in the example you shared, labels and scores are displayed on top of the image. However, for tasks like detection and inst seg, these elements should be displayed next to their corresponding bounding box or instance mask for better visual result.

Would it be possible to incorporate task-specific visualization logic into *Result(VisualizeMixin)? This could ensure a more flexible and consistent approach to rendering task-specific annotations.

For instance:

def _generate_simple(self, image: Image, result: VisualizeMixin) -> Image:
    return result.generate(image)

In this way, we can delegate the task-specific rendering to task result object, which would contain the appropriate logic for different tasks (e.g., detection, segmentation).