<reponame>apacha/maskrcnn-benchmark
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import cv2
import datetime
import logging
import time
import os

import torch
from tqdm import tqdm

from maskrcnn_benchmark.data.datasets.evaluation import evaluate
from ..utils.comm import is_main_process, get_world_size
from ..utils.comm import all_gather
from ..utils.comm import synchronize
import numpy as np


def compute_on_dataset(model, data_loader, device):
    model.eval()
    results_dict = {}
    cpu_device = torch.device("cpu")
    for i, batch in enumerate(tqdm(data_loader)):
        images, targets, image_ids = batch
        images = images.to(device)
        with torch.no_grad():
            output = model(images)
            output = [o.to(cpu_device) for o in output]
        results_dict.update(
            {img_id: result for img_id, result in zip(image_ids, output)}
        )
    return results_dict


def _accumulate_predictions_from_multiple_gpus(predictions_per_gpu):
    all_predictions = all_gather(predictions_per_gpu)
    if not is_main_process():
        return
    # merge the list of dicts
    predictions = {}
    for p in all_predictions:
        predictions.update(p)
    # convert a dict where the key is the index in a list
    image_ids = list(sorted(predictions.keys()))
    if len(image_ids) != image_ids[-1] + 1:
        logger = logging.getLogger("maskrcnn_benchmark.inference")
        logger.warning(
            "Number of images that were gathered from multiple processes is not "
            "a contiguous set. Some images might be missing from the evaluation"
        )

    # convert to a list
    predictions = [predictions[i] for i in image_ids]
    return predictions


def select_top_predictions(prediction, confidence_threshold=0.5):
    scores = prediction.get_field("scores")
    keep = torch.nonzero(scores > confidence_threshold).squeeze(1)
    prediction = prediction[keep]
    scores = prediction.get_field("scores")
    _, idx = scores.sort(0, descending=True)
    return prediction[idx]


def overlay_boxes(image, prediction):
    """
    Adds the predicted boxes on top of the image

    Arguments:
        image (np.ndarray): an image as returned by OpenCV
        prediction (BoxList): the result of the computation by the model.
            It should contain the field `labels`.
    """
    labels = prediction.get_field("labels")
    boxes = prediction.bbox

    colors = compute_colors_for_labels(labels).tolist()
    image = np.array(image).transpose((1, 2, 0))

    for box, color in zip(boxes, colors):
        box = box.to(torch.int64)
        top_left, bottom_right = box[:2].int().tolist(), box[2:].int().tolist()
        image = cv2.rectangle(image, tuple(top_left), tuple(bottom_right), tuple(color), 1)

    return image


def compute_colors_for_labels(labels):
    palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
    colors = labels[:, None] * palette
    colors = (colors % 255).numpy().astype("uint8")
    return colors


def visualize_results_in_images(dataset, predictions, output_folder):
    for data_point, prediction in zip(dataset, predictions):
        top_predictions = select_top_predictions(prediction)
        image = data_point[0]
        image = overlay_boxes(image, top_predictions)
        image_id = dataset.id_to_img_map[data_point[2]]
        cv2.imwrite(f"test-{image_id}.png", image)


def inference(
        model,
        data_loader,
        dataset_name,
        iou_types=("bbox",),
        box_only=False,
        device="cuda",
        expected_results=(),
        expected_results_sigma_tol=4,
        output_folder=None,
        visualize_results=True
):
    # convert to a torch.device for efficiency
    device = torch.device(device)
    num_devices = get_world_size()
    logger = logging.getLogger("maskrcnn_benchmark.inference")
    dataset = data_loader.dataset
    logger.info("Start evaluation on {} dataset({} images).".format(dataset_name, len(dataset)))
    start_time = time.time()
    predictions = compute_on_dataset(model, data_loader, device)
    # wait for all processes to complete before measuring the time
    synchronize()
    total_time = time.time() - start_time
    total_time_str = str(datetime.timedelta(seconds=total_time))
    logger.info(
        "Total inference time: {} ({} s / img per device, on {} devices)".format(
            total_time_str, total_time * num_devices / len(dataset), num_devices
        )
    )

    predictions = _accumulate_predictions_from_multiple_gpus(predictions)
    if not is_main_process():
        return

    if output_folder:
        torch.save(predictions, os.path.join(output_folder, "predictions.pth"))

    extra_args = dict(
        box_only=box_only,
        iou_types=iou_types,
        expected_results=expected_results,
        expected_results_sigma_tol=expected_results_sigma_tol,
    )

    if visualize_results:
        visualize_results_in_images(dataset, predictions, output_folder)

    return evaluate(dataset=dataset,
                    predictions=predictions,
                    output_folder=output_folder,
                    **extra_args)

# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
import cv2
import datetime
import logging
import time
import os

import torch
from tqdm import tqdm

from maskrcnn_benchmark.data.datasets.evaluation import evaluate
from ..utils.comm import is_main_process, get_world_size
from ..utils.comm import all_gather
from ..utils.comm import synchronize
import numpy as np


def compute_on_dataset(model, data_loader, device):
    model.eval()
    results_dict = {}
    cpu_device = torch.device("cpu")
    for i, batch in enumerate(tqdm(data_loader)):
        images, targets, image_ids = batch
        images = images.to(device)
        with torch.no_grad():
            output = model(images)
            output = [o.to(cpu_device) for o in output]
        results_dict.update(
            {img_id: result for img_id, result in zip(image_ids, output)}
        )
    return results_dict


def _accumulate_predictions_from_multiple_gpus(predictions_per_gpu):
    all_predictions = all_gather(predictions_per_gpu)
    if not is_main_process():
        return
    # merge the list of dicts
    predictions = {}
    for p in all_predictions:
        predictions.update(p)
    # convert a dict where the key is the index in a list
    image_ids = list(sorted(predictions.keys()))
    if len(image_ids) != image_ids[-1] + 1:
        logger = logging.getLogger("maskrcnn_benchmark.inference")
        logger.warning(
            "Number of images that were gathered from multiple processes is not "
            "a contiguous set. Some images might be missing from the evaluation"
        )

    # convert to a list
    predictions = [predictions[i] for i in image_ids]
    return predictions


def select_top_predictions(prediction, confidence_threshold=0.5):
    scores = prediction.get_field("scores")
    keep = torch.nonzero(scores > confidence_threshold).squeeze(1)
    prediction = prediction[keep]
    scores = prediction.get_field("scores")
    _, idx = scores.sort(0, descending=True)
    return prediction[idx]


def overlay_boxes(image, prediction):
    """
    Adds the predicted boxes on top of the image

    Arguments:
        image (np.ndarray): an image as returned by OpenCV
        prediction (BoxList): the result of the computation by the model.
            It should contain the field `labels`.
    """
    labels = prediction.get_field("labels")
    boxes = prediction.bbox

    colors = compute_colors_for_labels(labels).tolist()
    image = np.array(image).transpose((1, 2, 0))

    for box, color in zip(boxes, colors):
        box = box.to(torch.int64)
        top_left, bottom_right = box[:2].int().tolist(), box[2:].int().tolist()
        image = cv2.rectangle(image, tuple(top_left), tuple(bottom_right), tuple(color), 1)

    return image


def compute_colors_for_labels(labels):
    palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
    colors = labels[:, None] * palette
    colors = (colors % 255).numpy().astype("uint8")
    return colors


def visualize_results_in_images(dataset, predictions, output_folder):
    for data_point, prediction in zip(dataset, predictions):
        top_predictions = select_top_predictions(prediction)
        image = data_point[0]
        image = overlay_boxes(image, top_predictions)
        image_id = dataset.id_to_img_map[data_point[2]]
        cv2.imwrite(f"test-{image_id}.png", image)


def inference(
        model,
        data_loader,
        dataset_name,
        iou_types=("bbox",),
        box_only=False,
        device="cuda",
        expected_results=(),
        expected_results_sigma_tol=4,
        output_folder=None,
        visualize_results=True
):
    # convert to a torch.device for efficiency
    device = torch.device(device)
    num_devices = get_world_size()
    logger = logging.getLogger("maskrcnn_benchmark.inference")
    dataset = data_loader.dataset
    logger.info("Start evaluation on {} dataset({} images).".format(dataset_name, len(dataset)))
    start_time = time.time()
    predictions = compute_on_dataset(model, data_loader, device)
    # wait for all processes to complete before measuring the time
    synchronize()
    total_time = time.time() - start_time
    total_time_str = str(datetime.timedelta(seconds=total_time))
    logger.info(
        "Total inference time: {} ({} s / img per device, on {} devices)".format(
            total_time_str, total_time * num_devices / len(dataset), num_devices
        )
    )

    predictions = _accumulate_predictions_from_multiple_gpus(predictions)
    if not is_main_process():
        return

    if output_folder:
        torch.save(predictions, os.path.join(output_folder, "predictions.pth"))

    extra_args = dict(
        box_only=box_only,
        iou_types=iou_types,
        expected_results=expected_results,
        expected_results_sigma_tol=expected_results_sigma_tol,
    )

    if visualize_results:
        visualize_results_in_images(dataset, predictions, output_folder)

    return evaluate(dataset=dataset,
                    predictions=predictions,
                    output_folder=output_folder,
                    **extra_args)
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. # merge the list of dicts # convert a dict where the key is the index in a list # convert to a list Adds the predicted boxes on top of the image      Arguments:         image (np.ndarray): an image as returned by OpenCV         prediction (BoxList): the result of the computation by the model.             It should contain the field `labels`. # convert to a torch.device for efficiency # wait for all processes to complete before measuring the time