Tryag File Manager

//paip/script/weight/unused/unused_weightAnalyzer.py

import warnings
import os
import cv2
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from sklearn.cluster import DBSCAN
from sklearn.mixture import GaussianMixture as GMM
import time
from matplotlib.colors import ListedColormap
import seaborn as sns
import sys

if __name__ == '__main__':

working_dir = sys.argv[1] if len(sys.argv) > 1 else './'
    os.chdir(working_dir)
    default_img = 'new_img_6.jpg'
    warnings.filterwarnings("ignore")
    start_time = time.time()

# setup
    n_clusters = 3
    low_bound, upper_bound = 20000, 120000  # pixel number of each group
    data_dir = sys.argv[1] if len(sys.argv) > 1 else '../data/'
    output_dir = sys.argv[1] if len(sys.argv) > 1 else '../out/'
    print_time_flag = 0 if len(sys.argv) > 1 else 1
    img_name = os.path.basename(sys.argv[2]) if len(sys.argv) > 2 else default_img

# load image
    img = cv2.imread(os.path.join(data_dir, img_name))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

original_shape = img.shape

img2 = img.reshape((-1, 3))

# load palette
    cmap = ListedColormap(sns.color_palette("tab20"))
    # print(f'1/5...load image : {os.path.join(data_dir, img_name)}')
    # print(f'2/5...GMM model...')
    # GMM
    g_model = GMM(n_components=n_clusters, covariance_type='tied').fit(img2)

# choose 2 of 3 indexes ofr g_model results
    g_model_index = np.argsort(g_model.means_[:, 1])[-2:]

predicted = g_model.predict(img2)

img2_da = img2.copy()
    img_input = np.append(img2_da, [[x] for x in predicted], axis=1)
    img_data = pd.DataFrame(img_input.reshape(original_shape[0] * original_shape[1], 4),
                            columns=['r', 'g', 'b', 'gmm']).reset_index()

img_data_input = img_data[img_data.gmm.isin(g_model_index)]

img_data_input['x'] = img_data_input.index // original_shape[1]
    img_data_input['y'] = img_data_input.index % original_shape[1]

# print(f'2/5...GMM... : {time.time() - start_time}')
    # print(f'3/5...DB SCAN...')
    # DB SCAN
    model = DBSCAN(eps=1., min_samples=5, n_jobs=-1)
    pred = model.fit_predict(img_data_input[['x', 'y']])
    img_data_input['group'] = pred
    # print(f'3/5...DB SCAN... : {time.time() - start_time}')
    # pd.dataframe

result = pd.DataFrame(img_data_input.groupby('group')['group'].count()[2:])
    result.columns = ['cnt']

# group 별 좌표 구하기

grp_list = list(result[result.cnt > low_bound].index)

min_max_img_df = img_data_input[img_data_input.group.isin(grp_list)].groupby('group')[['x', 'y']].agg(
        min_x=('x', 'min'), max_x=('x', 'max'), min_y=('y', 'min'), max_y=('y', 'max'), cnt=('x', 'count')
    )

min_max_img_df['mat'] = (min_max_img_df.max_x - min_max_img_df.min_x) * (
            min_max_img_df.max_y - min_max_img_df.min_y)
    min_max_img_df['diffs'] = min_max_img_df.mat - min_max_img_df.cnt
    min_max_img_df['ratio'] = round(min_max_img_df.cnt / min_max_img_df.mat, 1)
    min_max_img_df = min_max_img_df.reset_index()

# pixel cnt... image boundary
    min_max_img_df = min_max_img_df[
        (min_max_img_df.cnt > low_bound) & (min_max_img_df.cnt < upper_bound) &  # pixel cnt filter
        (min_max_img_df.min_x > 0) &  # point filter below 4 lines
        (min_max_img_df.min_y > 0) &
        (min_max_img_df.max_x < original_shape[0] - 1) &
        (min_max_img_df.max_y < original_shape[1] - 1)]

# print(f'df : {min_max_img_df}')

# group candi
    final_group_candi = min_max_img_df.group.unique()

# replce to white color of each pixel which is in group candi
    # print(f'4/5...find contour and fill inner pixels')
    img_data_input = img_data_input[img_data_input.group.isin(min_max_img_df.group.unique())]
    '''
    img_cp = img.copy()

for x in img_data_input[['x', 'y', 'group']].iterrows():
        # print(x)
        x_ind, y_ind, grp_ind = x[1]['x'], x[1]['y'], x[1]['group']
        if grp_ind in grp_list:
            img_cp[x_ind, y_ind] = tuple(int(y * 255) for y in cmap(grp_ind % cmap.N)[:-1])
    '''
    # print(f'4/5...Image draw')
    img_cp = img.copy()

for x in min_max_img_df.iterrows():
        # print(x)
        # bound
        x_min, x_max, y_min, y_max, grp_ind = int(x[1]['min_x']), int(x[1]['max_x']), int(x[1]['min_y']), int(
            x[1]['max_y']), int(x[1]['group'])
        map_color = tuple(y * 255 for y in cmap(grp_ind % cmap.N)[:-1])
        inner_list = img_data_input[img_data_input.group == grp_ind][['x', 'y']]
        for xxx, yyy in inner_list.values.tolist():
            img_cp[xxx, yyy] = map_color

# find contours and fill them
    calib_list = []
    img_result = img.copy()
    for i, x in enumerate(final_group_candi):
        reg_color = tuple(int(y * 255) for y in cmap(x % cmap.N)[:-1])
        array_color = np.array(reg_color)

# cluster group dataframe
        info_df = min_max_img_df[min_max_img_df.group == x]

# image piece, slicing by info_df from origin-size image
        region_img = img_cp[int(info_df.min_x):int(info_df.max_x), int(info_df.min_y):int(info_df.max_y)].copy()

# 대상 클러스터의 유효한 pixel 을 white 로 변경. (contour 구하기용)
        x_ind, y_ind = region_img.shape[:2]
        for x_i in range(x_ind):
            for y_i in range(y_ind):
                if all(region_img[x_i, y_i] == array_color):
                    region_img[x_i, y_i] = (255, 255, 255)

gray = cv2.cvtColor(region_img, cv2.COLOR_RGB2GRAY)

ret, binary = cv2.threshold(gray, 254, 255, cv2.THRESH_BINARY)

# contour  구하기
        contours, hierarchy = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

# 구한 contour 내부 색을 cluster 색으로 칠함
        region_img_after = region_img.copy()
        region_img_after = cv2.fillPoly(region_img_after, pts=contours, color=reg_color)

# 해당 면적 크기 구하기
        white_list = [(x, y) for x in range(x_ind) for y in range(y_ind) if
                      all(region_img_after[x, y] == reg_color)]

calib_list.append(len(white_list))

# 원본 이미지에 생성한 이미지로 대체
        img_result[int(info_df.min_x):int(info_df.max_x), int(info_df.min_y):int(info_df.max_y)] = region_img_after

# 계산된 면적 column
    min_max_img_df['calib_cnt'] = calib_list

# draw rectangles, put text. ... on detected objects
    for i, x in enumerate(min_max_img_df.iterrows()):
        img_result = cv2.rectangle(img_result, (int(x[1]['min_y']), int(x[1]['min_x'])),
                                   (int(x[1]['max_y']), int(x[1]['max_x'])), (0, 255, 0), 3)
        img_result = cv2.putText(img_result, f'({i+1}) : {x[1]["calib_cnt"]}',
                                 (int(x[1]['min_y']), int(x[1]['max_x'] + 30)), cv2.FONT_HERSHEY_SIMPLEX, 1,
                                 (255, 0, 0), 3)
    # print(f'4/5...Image Draw : {time.time() - start_time}')
    # put text :: mean value
    # 결과 데이터 이미지내 삽입
    mean_value = round(min_max_img_df.calib_cnt.mean(), 1)
    img_result = cv2.putText(img_result,
                             f'mean : {min_max_img_df.calib_cnt.mean()}, median : {min_max_img_df.calib_cnt.median()}, n_detect : {len(min_max_img_df)}',
                             (5, original_shape[0] - 5), cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 0, 0), 2)

# plot image
    f, (ax1, ax2) = plt.subplots(1, 2, figsize=(24, 9))
    f.tight_layout()
    ax1.imshow(img)
    ax1.set_title('Original Image', fontsize=50)
    ax2.imshow(img_result)
    ax2.set_title('clustered Image', fontsize=50)
    plt.subplots_adjust(left=0., right=1, top=0.9, bottom=0.)

# 이미지 저장 : result_<<원 이미지 파일 이름>>
    f.savefig(os.path.join(output_dir, 'result_' + img_name))

# console out
    print(
        f'mean : {min_max_img_df.calib_cnt.mean()}, median : {min_max_img_df.calib_cnt.median()}, n_detect : {len(min_max_img_df)}')

# test용 print
    if print_time_flag : print(f'exec duration : {time.time() - start_time}')