import torch


def get_bbox(image_id, target, device='cpu', pred=False):
    """
        image_id (int): image id from annotation.
        target (dict): target data from annotation or result of model inference.
        device (string): use CUDA of CPU.
        pred (bool): if True, result include confidence score of bbox.
            if False, result include only bboxes and labels.
            
        
    """
    
    box_num = target['boxes'].shape[0]
    
    image_id = torch.tensor(image_id).broadcast_to(box_num, 1)
    boxes = target['boxes'].reshape(box_num, 4).to(device)
    labels = target['labels'].reshape(box_num, 1).to(device)
    
    if pred:
        scores = target['scores'].reshape(box_num, 1).to(device)
        bbox_list = torch.cat([image_id, labels, scores, boxes], dim=1)
        return bbox_list
    
    bbox_list = torch.cat([image_id, labels, boxes], dim=1)
    
    return bbox_list