sex_age_plot.py

import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

from matplotlib.pyplot import imread

def sex_age_plot(cluster: int, save=False):
    # 이미지 불러오기
    female_images = [imread(f'../images/ages icon/F{i}.png') for i in range(10,80,10)]
    male_images = [imread(f'../images/ages icon/M{i}.png') for i in range(10,80,10)]    
    
    images = female_images + male_images
    
    np.random.seed(223) # it's my birthday
    sex_age = [f'{sex}_{age}' for sex in ['F','M'] for age in range(10,80,10)]
    cluster_df = pd.DataFrame({'sex_age':sex_age, 'k':np.random.randint(low=0, high=5, size=len(sex_age))})

    cluster_idx = cluster_df[cluster_df['k']==cluster].index

    # 이미지 삽입
    plt.figure(figsize=(11,6))
    for i, image in enumerate(images):
        # 군집 이외의 ID는 불투명 상태로 변환
        alpha = 1 if i in cluster_idx else 0.3

        height, width, space = 40, 20, 2
        left = width*i + space*i
        right = left + width
        bottom = 0

        if i > 6:
            height += 40
            bottom = 40
            left = width*(i-7) + space*(i-7)
            right = left + width
            
        # 색상 - 남: 파랑, 여: 빨강
        if i > 6:
            # 남자
            image[image!=0] = 255
            image[np.where((image==[0,0,0]).all(axis=2))] = [0,0,255]
        else:
            # 여자
            image[image!=0] = 255
            image[np.where((image==[0,0,0]).all(axis=2))] = [255,0,0]
            
        image /= 255.
            
        plt.imshow(image, extent=[left, right, bottom, height], alpha=alpha)

    plt.xlim(0,152)
    plt.ylim(0,80)
    plt.axis('off')

    if save:
        plt.tight_layout()
        plt.savefig('../images/sample/sex_age_example.jpg', dip=500)
        
    plt.show()