utils_vis.py

Go to the documentation of this file.

import numpy as np
# import mmcv
import copy
import cv2
from tqdm import tqdm
 
import smutils.utils_os as utils
 
# from mmpose.apis import inference_top_down_pose_model, init_pose_model, vis_pose_result
 
def get_palette(n):
    """
    설정한 수량만큼 컬러 테이블을 생성하는 기능
    args:
        n (int): 생성하려는 색상 수량
    return (list[tuple]): 컬러 테이블
    """
    state = np.random.get_state()
    np.random.seed(42) #고정된 컬러 테이블 출력을 위해 셋팅
    palette = np.random.randint(0, 256, size=(n, 3))
    np.random.set_state(state) #이후에 랜덤성 관련 함수를 사용할 수도 있기 때문에 원래의 값으로 셋팅
    colors = [tuple(c) for c in palette]
    return colors
 
def vis_instance_segmentation(image, masks=None, bboxes=None, colors=None):
    """
    Instance Segmentation 결과 시각화를 위한 기능
    박스 또는 마스크 둘중 하나 이상은 입력되어야함
    args:
        image (np.array): 이미지
        masks (list[np.array]): Instance Mask, 각 마스크의 shape은 이미지의 크기와 동일함
        bboxes (np.array): 검출 객체의 박스, 순서는 마스크와 매칭되어야함
        colors (list[tuple]): 컬러 테이블
    return (np.array): Instance Segmentation 결과가 시각화된 이미지
        TODO: 이미지를 리턴하지 않아도 됨, 리턴 제거
    """
 
    #컬러 테이블 설정
    if colors is None:
        if masks is not None:
            n = masks.shape[0]
        else:
            n = bboxes.shape[0]
        colors = get_palette(n)
 
    #마스크가 입력되었을 경우
    if masks is not None:
        mask_img = np.zeros_like(image) #마스크 이미지 초기화
        
        #마스크 이미지 생성
        for mask, color in zip(reversed(masks), colors):
            mask_img[mask] = color
        
        #컬러 영상에 오버랩
        mask = mask_img > 0
        image[mask] = image[mask]*0.3 + mask_img[mask]*0.7
 
    #박스가 입력되었을 경우 박스 그리기
    if bboxes is not None:
        for bbox, color in zip(reversed(bboxes), colors):
            image = draw_single_bbox_and_label(image, bbox, f'{bbox[-1]:.3f}', color, 2, color, 2, 0.7, box_type='xyxy')
 
    return image
 
def vis_pose_coco_skeleton(image, pose_results):
    """
    coco 스켈레톤을 시각화하는 기능
    args:
        image (np.array): 이미지
        pose_results (list[dict] or dict): 스켈레톤 정보를 담고 'keypoints'가 포함된 데이터 리스트
    return (np.array): coco 스켈레톤이 시각화된 이미지
    """
 
    #하나의 데이터일 경우 리스트로 감싸기
    if not isinstance(pose_results, list):
        pose_results = [pose_results]
 
    #시각화
    for pose in pose_results:
        kps = pose['keypoints'] if pose['keypoints'].ndim == 2 else pose['keypoints'][0]    
        # kps[:, :2] += pose['bbox'][:2] #TODO: 크롭 이미지기반의 좌표일 경우 원본이미지 좌표로 이동, 옵션 처리 필요
        image = vis_pose_coco_skeleton_one_person(image, kps)
 
    return image
 
def vis_pose_coco_skeleton_one_person(image, keypoints):
    """
    한명에 대한 coco 스켈레톤을 시각화 하는 기능
    args:
        image (np.array): 이미지
        keypoints (np.array): 한명에 대한 스켈레톤 좌표, shape (17, 3)
    return (np.array): 스켈레톤이 시각화된 이미지
    """
    # COCO ordered keypoints:
    # 0:nose  1:left_eye  2:right_eye  3:left_ear  4:right_ear  5:left_shoulder
    # 6:right_shoulder  7:left_elbow  8:right_elbow  9:left_wrist  10:right_wrist
    # 11:left_hip  12:right_hip  13:left_knee  14:right_knee  15:left_ankle  16:right_ankle
 
    # define connections between keypoints for COCO dataset
    coco_pairs = [
        (0, 1), (0, 2), (1, 3), (2, 4),  # head
        (5, 7), (7, 9),  # left arm
        (6, 8), (8, 10),  # right arm
        (5, 6), (5, 11), (6, 12), (11, 12),  # body
        (11, 13), (13, 15),  # left leg
        (12, 14), (14, 16)  # right leg
    ]
 
    # define colors for different parts of the body
    colors = {
        'face': (255, 210, 127),  # blue
        'left_body': (0, 165, 255),  # orange
        'right_body': (0, 255, 0),  # green
        'body': (255, 210, 127)  # blue
    }
 
    #스켈레톤에 대한 컬러맵
    keypoint_colors = [colors['face']] * 5 + [colors['left_body'], colors['right_body']] * 6
 
    #관절에 대한 컬러맵
    pair_colors = [colors['face']] * 4 + [colors['left_body']] * 2 + [colors['right_body']] * 2 + \
                [colors['body']] * 4 + [colors['left_body']] * 2 + [colors['right_body']] * 2
 
 
    #관절 그리기
    for i, pair in enumerate(coco_pairs):
        kp1, kp2 = pair
 
        #관절을 연결하는 두개의 스켈레톤의 스코어가 모두 0.5 이상일경우 그리기
        if keypoints[kp1, 2] > 0.5 and keypoints[kp2, 2] > 0.5: 
            cv2.line(image, (int(keypoints[kp1, 0]), int(keypoints[kp1, 1])),
                    (int(keypoints[kp2, 0]), int(keypoints[kp2, 1])), pair_colors[i], 2)
 
    #스켈레톤 그리기
    for i, keypoint in enumerate(keypoints):
        x, y, score = keypoint
 
        #스켈레톤의 스코어가 0.5 이상이면 그리기
        if score > 0.5:
            cv2.circle(image, (int(x), int(y)), 5, keypoint_colors[i], thickness=-1, lineType=cv2.FILLED)
 
    return image
 
#TODO: 박스 그리는 함수들을 일반화 시킬 필요가 있음
def draw_single_bbox_and_label(img, bbox, label, box_color, box_thk, txt_color, txt_thk, txt_scale, box_type='xywh'):
    """
    이미지에 박스 한개와 텍스트를 출력하는 기능
    텍스트의 위치는 박스의 상단, 상단에 출력하지 못할것 같다고 판단되면 박스 하단에 출력
    args:
        img (np.array): 이미지
        bbox (np.array): 박스, shape (4 or 5)
        label (str): 출력 문구
        box_color (tuple): 박스 컬러
        box_thk (int): 박스를 그리는 선의 두께
        txt_color (tuple): 텍스트의 색상
        txt_thk (int): 텍스트 선의 두께
        txt_scale (float): 텍스트의 크기
        box_type (str): 박스 구조, 'xyxy', 'xywh' 지원
    return (np.array): 박스 및 텍스트가 그려진 이미지
    """
 
    #텍스트 출력 위치
    txt_xy = (int(bbox[0])-30, int(bbox[1]-20))
    if int(bbox[1]-20) < 40:
        txt_xy = (int(bbox[0]-30), int(bbox[1]+bbox[3]+25))
    
    #입력 데이터를 int로 변환후 튜플로 묶기
    box_color = tuple(map(int, box_color))
    txt_color = tuple(map(int, txt_color))
    
    #박스 그리기
    if box_type == 'xywh':
        img = cv2.rectangle(img, (int(bbox[0]), int(bbox[1])), (int(bbox[0]+bbox[2]), int(bbox[1]+bbox[3])), box_color, box_thk)
    elif box_type == 'xyxy':
        img = cv2.rectangle(img, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), box_color, box_thk)
 
    #텍스트 출력
    cv2.putText(img, label, txt_xy, cv2.FONT_HERSHEY_DUPLEX, txt_scale, txt_color, txt_thk, 1)
 
    return img
 
def draw_bboxes_list(image, bboxes, box_color, box_thk, box_type, labels=None):
    """
    리스트 또는 np.array 타입의 박스를 이미지에 그리는 기능
    args:
        image (np.array): 이미지
        bboxes (list[np.array], np.array): 박스
        box_color (tuple): 박스 컬러
        box_thk (int): 박스를 그리는 선의 두께
        box_type (str): 박스 구조, 'xyxy', 'xywh' 지원
        labels (str): 출력 문구
    return (np.array): 박스 및 텍스트가 그려진 이미지
    """
 
    #출력 문구가 없을경우 초기화
    if labels == None:
        labels = ['']*len(bboxes)
 
    #박스 그리기
    for bbox, label in zip(bboxes, labels):
         #출력 문구를 추적 아이디로 설정하고 박스 셋팅
        if label == '' and bbox.shape[0] == 6:
            label = str(int(bbox[0]))
            bbox = bbox[1:]
 
        #그리기    
        image = draw_single_bbox_and_label(image, bbox, label, box_color, box_thk, (255, 255, 255), 2, 1.5, box_type=box_type)
    return image
 
def draw_bboxes_dict(image, bboxes, box_color, box_thk, box_type):
    """
    딕셔너리 형태로 입력된 박스를 그리는 기능
    키는 추적 아이디, value는 박스
    args:
        image (np.array): 이미지
        bboxes (dict[np.array]): 박스
        box_color (tuple): 박스 컬러
        box_thk (int): 박스를 그리는 선의 두께
        box_type (str): 박스 구조, 'xyxy', 'xywh' 지원
    return (np.array): 박스가 그려진 이미지
    """
    for track_id, bbox in bboxes.items():
        label = f'{track_id}'
        image = draw_single_bbox_and_label(image, bbox, label, box_color, box_thk, (255, 255, 255), 2, 1.5, box_type=box_type)
    return image
 
def draw_and_save_bboxes(image, save_path, bboxes, box_color, box_thk, box_type='xywh'):
    """
    박스를 그리고 박스가 그려진 이미지를 저장하는 기능
    args:
        image (str or np.array): 이미지 경로 또는 이미지
        save_path (str): 이미지 저장 경로
        bboxes (np.array or list[np.array] or dict): 박스
        box_color (tuple): 박스 컬러
        box_thk (int): 박스를 그리는 선의 두께
        box_type (str): 박스 구조, 'xyxy', 'xywh' 지원
    """
 
    #이미지의 입력이 경로인 경우 로드
    if isinstance(image, str):
        image = cv2.imread(image)
    
    #그리기
    if isinstance(bboxes, dict):
        image = draw_bboxes_dict(image, bboxes, box_color, box_thk, box_type)
    else:
        image = draw_bboxes_list(image, bboxes, box_color, box_thk, box_type)
 
    #저장
    cv2.imwrite(save_path, image)
 
def draw_and_save_single_bbox(image, save_path, bbox, label, box_color, box_thk, box_type='xywh'):
    """
    박스 한개를 그리고 박스가 그려진 이미지를 저장하는 기능
    args:
        image (str or np.array): 이미지 경로 또는 이미지
        save_path (str): 이미지 저장 경로
        bbox (np.array): 박스, shape (4 or 5 or 6)
        label (str): 출력 문구
        box_color (tuple): 박스 컬러
        box_thk (int): 박스를 그리는 선의 두께
        box_type (str): 박스 구조, 'xyxy', 'xywh' 지원
    """
 
    #이미지의 입력이 경로인 경우 로드
    if isinstance(image, str):
        image = cv2.imread(image)
    
    #박스에 추적 아이디가 포함된 경우
    if bbox.shape[0] == 6:
        bbox = bbox[1:]
    
    #그리기
    image = draw_single_bbox_and_label(image, bbox, label, box_color, box_thk, (255, 255, 255), 2, 1.5, box_type=box_type)
    
    #저장
    cv2.imwrite(save_path, image)        
 
def make_video(folder_or_paths, save_path, fps=30, i_size='half'):
    """
    이미지를 기반으로 동영상을 생성하는 기능
    TODO: utils_video.py의 make_video를 사용하도록 함, 제거 필요
    args:
        folder_or_paths (str or list[str]): 동영상을 만들 이미지가 저장된 폴더 또는 이미지 경로 리스트
        save_path (str): 동영상을 저장할 경로
        fps (int): 생성할 동영상의 프레임 속도
        i_size (str): 이미지의 크기를 절반으로 줄여서 저장할지 여부, 'half'입력했을 경우 영상을 절반 크기로 생성함
    """
 
    print('Warning: Please use make_video in utils_video.py.')
 
    #이미지 파일경로 설정
    if isinstance(folder_or_paths, str):
        _, path_list = utils.search_file(folder_or_paths, '.jpg')
    else:
        path_list = folder_or_paths
 
    #첫 이미지 읽어서 이미지 크기 알아내기
    img = cv2.imread(path_list[0])
    h, w, _ = img.shape
 
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
 
    #이미지 크기 설정    
    if i_size == 'half':
        h = int(h/2)
        w = int(w/2)
    
    #비디오 작성자 생성
    video_writer = cv2.VideoWriter(save_path, fourcc, fps, (w, h))
 
    #동영상 만들기
    N = len(path_list)
    for i in tqdm(range(N)):
        path = path_list[i]
        img = cv2.imread(path) #이미지 읽기
 
        if i_size == 'half': #이미지 사이즈 조절
            img = cv2.resize(img, dsize=(w, h), interpolation=cv2.INTER_AREA)
 
        video_writer.write(img)
 
    print()
 
    video_writer.release()