Tryag File Manager

//paip/script/distribution_weight/img_distribution_for_weight.py

# Copyright (c) 2020~2024 PAIPTREE SmartFarm, Inc. All Rights Reserved.
# Modifier: 박정훈 2024-02-21
# bugfix : 암전 판단시 이미지를 중심 기준으로 전체 크기 * 0.5배 사이즈로 크롭하고 있는데
# bugfix : 암전 판단시 width는 1배, height는 0.9배 사이즈 크롭으로 변경

'''
Created on Mon Aug 30 2023
@author: Park Jeonghoon
'''

from datetime import datetime
import getopt, sys, cv2, os, warnings, glob
from matplotlib import patches, pyplot as plt
from matplotlib.ticker import MultipleLocator
import numpy as np
import pandas as pd 
from itertools import product
from collections import defaultdict
from sklearn.mixture import GaussianMixture as GMM

warnings.filterwarnings('ignore')
HOME_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), os.pardir)
MODULE_DIR = 'util'
sys.path.append(os.path.join(HOME_PATH, MODULE_DIR))
from PyDBconnector import PyDBconnector
from logs import paiplog

## 감지 threshold to alert
REGIONTHRESHOLD=.4
SUCCESS_MSG = 'success'
FAILED_MSG = 'fail'
BLACKOUT_MSG = 'blackout'

RESULT_FILE_SUFFIX = '_distributionResult'
IS_GATEWAY = True

########### 암전 판단 로직 ################################
def enhance_red_channel(image):
    source = image.astype(np.float32)
    blue, green, red = cv2.split(source)
    enhanced_red = 2 * red - green - blue

return enhanced_red

def get_h_channel_median(image):
    hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
    hue_channel, _, _ = cv2.split(hsv_image)
    h_channel_median = np.median(hue_channel)

return h_channel_median

def is_blackout(image_path): # 암전 판단 로직
    result = False
    ratio_width = 1
    ratio_height = 0.9
    
    img = cv2.imread(image_path)
    height, width, _ = img.shape
    center_x, center_y = width // 2, height // 2

crop_width, crop_height = int(width * ratio_width), int(height * ratio_height)
    crop_half_width, crop_half_height = crop_width // 2, crop_height // 2

# x1 = center_x - crop_half_width
    # y1 = center_y - crop_half_height
    # x2 = center_x + crop_half_width
    # y2 = center_y + crop_half_height
    # 위와 같은 코드이지만 성능은 아래가 더 나음
    x1 = max(0, center_x - crop_half_width)
    y1 = max(0, center_y - crop_half_height)
    x2 = min(width, center_x + crop_half_width)
    y2 = min(height, center_y + crop_half_height)

cropped_image = img[y1:y2, x1:x2]
    
    h_median = get_h_channel_median(cropped_image)

if h_median >= 145: # IR 모드(X) 야간 : 145 이상
        result = True   
    elif h_median == 0: # IR 모드(O) 야간 : 0
        enhanced_red = enhance_red_channel(cropped_image)
        total_sum = np.sum(enhanced_red)
        if total_sum == 0:
            result = True
    
    return result

#############################################################

def createFolder(directory):   
    try:
        if not os.path.exists(directory):
            os.makedirs(directory)
    except OSError:
        print ('Error: Creating directory. ' +  directory)

def load_image(imgName) : 
    img = cv2.imread(imgName)
    if img is None : raise Exception('image is not loaded')
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
    
    return img, gray

def applyMask(img_id : str, img, isGray :bool = True) : 
    img_cMap = np.array(cMap[img_id])
    mask = np.zeros(img.shape[:2], np.uint8)
    # maskVal = 255 if isGray else (255, 255, 255)

if len(img_cMap) <= 3 : 
        _ = cv2.drawContours(mask, [*img_cMap], -1, (255, 255, 255), -1, cv2.LINE_AA)
    else : 
        _ = cv2.drawContours(mask, [img_cMap], -1, (255, 255, 255), -1, cv2.LINE_AA)
    dst = cv2.bitwise_and(img, img, mask=mask)
    return dst, np.sum(np.where(mask > 0 , 1, 0)), mask

@paiplog
def getclusteredImage(imgDir, isPlot=False, isDir=True) : 
    
    weightMap = [x/divNum[1]/2+ .5 for x in range(1,divNum[1]+1)] 
    if isDir : 
        images = glob.glob(imgDir)
        dirName = OUTPUT_DIR
    else : 
        images = [imgDir]
        dirName = os.path.dirname(imgDir)
    results_df = pd.DataFrame(columns=['fileName','file_time','numOfInterest','interestRatio'])

for idx, imgName in enumerate(images) : 
        if IS_GATEWAY :
            #imgName : /gate/data/filesImageReal/20231108/H02_CT03,1_20231108122102_farm_image_real_2295ae414e7f.jpg
            img_id = imgName.split('_')[1]
            dateStr = imgName.split('_')[2]
        else : 
            #imgName : /paip/data/filesImageReal/FA0001/20231108/FA0001_GW00_H02_CT03,1_20231108122102_farm_image_real_2295ae414e7f.jpg
            img_id = imgName.split('_')[3]
            dateStr = imgName.split('_')[4]

img = cv2.imread(imgName)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img, sumMask, mask = applyMask(img_id, img, isGray=False)
        original_shape = img.shape
        img2 = img.reshape((-1, 3))
        n_clusters = 4
        g_model = GMM(n_components=n_clusters, covariance_type='tied', random_state=0).fit(img2)
        g_model_index = np.argsort(g_model.means_[:,1])[-(n_clusters-2):]
        
        predicted = g_model.predict(img2)
        
        color_map_0 = {x: 255 if x in g_model_index else 0 for x in np.unique(predicted)}
        predicted_color_0 = np.array([color_map_0.get(x) for x in predicted])
        result_image_0 = predicted_color_0.reshape(tuple(original_shape[:2]))
        sumRatio = []

# plot
        if isPlot : 
            _, ax = plt.subplots(figsize=(8,4))
            _ = ax.imshow(result_image_0)

ax.xaxis.set_major_locator(MultipleLocator(grid_xsize))
            ax.xaxis.set_major_formatter('{x:.0f}')
            ax.yaxis.set_major_locator(MultipleLocator(grid_ysize))
            ax.yaxis.set_major_formatter('{x:.0f}')

except_list = exceptMap[img_id]
            for idx, (x, y) in enumerate(ziped_xy) : 
                mask_sum = np.sum(np.where(mask[y:y+grid_ysize,x:x+grid_xsize] > 0 , 1, 0))
                img_sum = np.sum(np.where(mask[y:y+grid_ysize,x:x+grid_xsize] > 0 , 1, 0) & np.where(result_image_0[y:y+grid_ysize,x:x+grid_xsize] > 0, 1, 0))
                ratioRegion = np.round(img_sum/mask_sum,1) if mask_sum > 0 else 0.
                # sumRatio.append(np.round(weightMap[int(idx%divNum[1])] * ratioRegion,2))
                sumRatio.append(np.round(ratioRegion,2))
                if idx not in except_list : 
                    if ratioRegion <= regionThreshold :             
                        rect=patches.Rectangle((x,y),grid_xsize,grid_ysize, 
                            #fill=False,
                            alpha=0.5,
                        facecolor="red")
                        ax.add_patch(rect)

else :
                        pass

else :
                    pass

_ = plt.grid()
            plt.tight_layout()

createFolder(dirName)

plt.savefig(f"{dirName}/{os.path.splitext(os.path.basename(imgName))[0]}{RESULT_FILE_SUFFIX}.jpg")

# plt.show()
        else : 
            for idx, (x, y) in enumerate(ziped_xy) : 
                mask_sum = np.sum(np.where(mask[y:y+grid_ysize,x:x+grid_xsize] > 0 , 1, 0))
                img_sum = np.sum(np.where(mask[y:y+grid_ysize,x:x+grid_xsize] > 0 , 1, 0) & np.where(result_image_0[y:y+grid_ysize,x:x+grid_xsize] > 0, 1, 0))
                ratioRegion = np.round(img_sum/mask_sum,1) if mask_sum > 0 else 0.
                # sumRatio.append(np.round(weightMap[int(idx%divNum[1])] * ratioRegion,2))
                sumRatio.append(np.round(ratioRegion,2))
            
        results_df.loc[len(results_df)] = [imgName, dateStr, sumMask, sumRatio]
        
    return results_df
    
def getDBData(dbConn, sql_str) :
    try :
        return dbConn.select_from_db(sql_str)
    except :
        return None

@paiplog
def insertDatabase(dbConn, data_df, isInsert=True) :
    
    try:
        create_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        
        if isInsert : # insert new InDate
            
            for pdRow in data_df.iterrows():
                aFileName = os.path.splitext(pdRow[1]['org_file_name'])
                tid= aFileName[0].split('_')[-1]
                if IS_GATEWAY :
                    insert_string = f'''insert into tbl_image_analysis_distribution (create_time, house_id, module_id, data_type, file_time, org_file_name, result_file_name, cluster_num_of_interest, cluster_interest_ratio_list, cluster_prediction_status) 
                                    values('{create_time}', '{pdRow[1]['house_id']}', '{pdRow[1]['module_id']}', 'real', '{pdRow[1]['file_time']}', '{pdRow[1]['org_file_name']}', 
                                    '{aFileName[0]+RESULT_FILE_SUFFIX+aFileName[1]}', 
                                    {pdRow[1]['numOfInterest']}, '{pdRow[1]['interestRatio']}', 'success');'''
                else :
                    insert_string = f'''insert into tbl_farm_image_analysis_distribution (create_time, farm_id, gateway_id, house_id, module_id, data_type, file_time, org_file_name, result_file_name, cluster_num_of_interest, cluster_interest_ratio_list, cluster_prediction_status) 
                                    values('{create_time}', '{pdRow[1]['farm_id']}', '{pdRow[1]['gateway_id']}', '{pdRow[1]['house_id']}', '{pdRow[1]['module_id']}', 'real', '{pdRow[1]['file_time']}', '{pdRow[1]['org_file_name']}', 
                                    '{aFileName[0]+RESULT_FILE_SUFFIX+aFileName[1]}', 
                                    {pdRow[1]['numOfInterest']}, '{pdRow[1]['interestRatio']}', 'success');'''
                dbConn.insert_to_db(insert_string)

return True
        else : # update to modify OutDate
            pass
        
    except Exception as e:
        raise Exception('insert query error!')

@paiplog
def main(argv): 
    FILE_NAME = argv[0]
    imgHeight, imgWidth = 1080, 1920 
    divNum = [24, 24]
    grid_xsize, grid_ysize = int(imgWidth / divNum[0]), int(imgHeight / divNum[1])
    # print(f"grid size : {grid_xsize, grid_ysize}")
    grid_x_all = [x for x in range(0,imgWidth,int(grid_xsize))]
    grid_y_all = [x for x in range(0,imgHeight,int(grid_ysize))]
    ziped_xy = list(product(grid_x_all, grid_y_all))
    regionThreshold = REGIONTHRESHOLD
    FileName=None
    try:
        # opts: getopt 옵션에 따라 파싱 ex) [('-i', 'myinstancce1')]
        # etc_args: getopt 옵션 이외에 입력된 일반 Argument
        # argv 첫번째(index:0)는 파일명, 두번째(index:1)부터 Arguments
        opts, etc_args = getopt.getopt(argv[1:], \
                                "h:i:w:d:t:f:", ["help","imgHeight=","imgWidth=","divNum=","regionThreshold=",'FileName='])

except getopt.GetoptError: # 옵션지정이 올바르지 않은 경우
        print(FILE_NAME, '-ih <imgHeight> -iw <imgWidth> -d <divNum> -t <regionThreshold>')
        sys.exit(2)

for opt, arg in opts: # 옵션이 파싱된 경우
        if opt in ("-h", "--help"): # HELP 요청인 경우 사용법 출력
            print(FILE_NAME, '-ih <imgHeight> -iw <imgWidth> -d <divNum>::[width_num, height_num] -t <regionThreshold> -f <FileName>')
            sys.exit()
        elif opt in ("-i", "--imgHeight"):
            imgHeight = int(arg)

elif opt in ("-w", "--imgWidth"):
            imgWidth = int(arg)

elif opt in ("-d", "--divNum"):
            divNum = eval(arg)
        elif opt in ("-t", "--regionThreshold"):
            regionThreshold = float(arg)
        elif opt in ("-f", "--FileName"):
            FileName = arg

return (imgHeight, imgWidth, divNum, grid_xsize, grid_ysize, grid_x_all, grid_y_all, ziped_xy, regionThreshold, FileName)

if __name__ == '__main__' : 
    try : 
        imgHeight, imgWidth, divNum, grid_xsize, grid_ysize, grid_x_all, grid_y_all, ziped_xy, regionThreshold, FileName = main(sys.argv)
        
        if FileName is None:
            raise Exception('Not Exist FileName!!')
        
        if os.path.basename(FileName).startswith('FA'): 
            IS_GATEWAY = False
        
        if is_blackout(FileName) == False: # 암전 상황이 아니면......
        
            dbConn = PyDBconnector()
            
            cMap = defaultdict(lambda : [[0,0],[imgWidth,0],[imgWidth,imgHeight],[0,imgHeight]])
            exceptMap = defaultdict(lambda : [0])
            
            results_df = getclusteredImage(FileName, isPlot=True, isDir=False)
            
            # create columns
            results_df['org_file_name'] = [os.path.basename(x) for x in results_df.fileName]
            results_df['farm_id']    = [x[0] if x[0].startswith('FA') else ''    for x in results_df.org_file_name.str.split('_')]
            results_df['gateway_id'] = [x[1] if x[0].startswith('FA') else ''    for x in results_df.org_file_name.str.split('_')]
            results_df['house_id']   = [x[2] if x[0].startswith('FA') else x[0]  for x in results_df.org_file_name.str.split('_')]
            results_df['module_id']  = [x[3] if x[0].startswith('FA') else x[1]  for x in results_df.org_file_name.str.split('_')]

if len(results_df) > 0 : 
                insertDatabase(dbConn, results_df, isInsert=True)
                print(SUCCESS_MSG)

else: # 암전 상황이면......
            print(BLACKOUT_MSG)

except : 
        print(FAILED_MSG)