292 lines
10 KiB
Python
292 lines
10 KiB
Python
import cv2
|
|
import pandas as pd
|
|
from sklearn.cluster import MiniBatchKMeans
|
|
import numpy as np
|
|
#import matplotlib.pyplot as plt
|
|
from tqdm import tqdm
|
|
from skimage import color
|
|
from .ld_convertor import skimage_rgb2lab, df2rgba, rgba2df, hsv2df, rgb2df, mask2df
|
|
#from .ld_utils import img_plot
|
|
from .bg_remover import get_foreground
|
|
|
|
def calc_ciede(mean_list, cls_list):
|
|
cls_no = []
|
|
tgt_no = []
|
|
ciede_list = []
|
|
for i in range(len(mean_list)):
|
|
img_1 = np.array(mean_list[i][:3])
|
|
for j in range(len(mean_list)):
|
|
if i == j:
|
|
continue
|
|
img_2 = np.array(mean_list[j][:3])
|
|
ciede = color.deltaE_ciede2000(skimage_rgb2lab(img_1), skimage_rgb2lab(img_2))[0][0]
|
|
ciede_list.append(ciede)
|
|
cls_no.append(cls_list[i])
|
|
tgt_no.append(cls_list[j])
|
|
ciede_df = pd.DataFrame({"cls_no": cls_no, "tgt_no": tgt_no, "ciede2000": ciede_list})
|
|
return ciede_df
|
|
|
|
def get_mask(df, cls_no):
|
|
mask = df.copy()
|
|
mask.loc[df["label"] != cls_no, ["r","g","b"]] = 0
|
|
mask.loc[df["label"] == cls_no, ["r","g","b"]] = 255
|
|
mask = cv2.cvtColor(df2rgba(mask).astype(np.uint8), cv2.COLOR_RGBA2GRAY)
|
|
return mask
|
|
|
|
def fill_mean_color(img_df, mask):
|
|
df_img = df2rgba(img_df).astype(np.uint8)
|
|
if len(df_img.shape) == 3:
|
|
mask = np.repeat(mask[:, :, np.newaxis], df_img.shape[-1], axis=-1)
|
|
masked_img = np.where(mask == 0, 0, df_img)
|
|
mean = np.mean(masked_img[mask != 0].reshape(-1, df_img.shape[-1]), axis=0)
|
|
|
|
img_df["r"] = mean[0]
|
|
img_df["g"] = mean[1]
|
|
img_df["b"] = mean[2]
|
|
|
|
return img_df, mean
|
|
|
|
def get_blur_cls(img, cls, size):
|
|
blur_img = cv2.blur(img, (size, size))
|
|
blur_df = rgba2df(blur_img)
|
|
blur_df["label"] = cls
|
|
img_list = []
|
|
mean_list = []
|
|
cls_list = list(cls.unique())
|
|
for cls_no in tqdm(cls_list):
|
|
mask = get_mask(blur_df, cls_no)
|
|
img_df = blur_df.copy()
|
|
img_df.loc[blur_df["label"] != cls_no, ["a"]] = 0
|
|
img_df, mean = fill_mean_color(img_df, mask)
|
|
df_img = df2rgba(img_df).astype(np.uint8)
|
|
img_list.append(df_img)
|
|
mean_list.append(mean)
|
|
return img_list, mean_list, cls_list
|
|
|
|
def get_cls_update(ciede_df, df, threshold):
|
|
set_list = [frozenset({cls, tgt}) for cls, tgt in ciede_df[ciede_df['ciede2000'] < threshold][['cls_no', 'tgt_no']].to_numpy()]
|
|
merge_set = []
|
|
while set_list:
|
|
set_a = set_list.pop()
|
|
merged = False
|
|
for i, set_b in enumerate(merge_set):
|
|
if set_a & set_b:
|
|
merge_set[i] |= set_a
|
|
merged = True
|
|
break
|
|
if not merged:
|
|
merge_set.append(set_a)
|
|
merge_dict = {}
|
|
for merge in merge_set:
|
|
cls_counts = {cls: len(df[df['label'] == cls]) for cls in merge}
|
|
max_cls = max(cls_counts, key=cls_counts.get)
|
|
for cls in merge:
|
|
merge_dict[cls] = max_cls
|
|
return merge_dict
|
|
|
|
|
|
def get_color_dict(mean_list, cls_list):
|
|
color_dict = {}
|
|
for idx, mean in enumerate(mean_list):
|
|
color_dict.update({cls_list[idx]:{"r":mean[0],"g":mean[1],"b":mean[2], }})
|
|
return color_dict
|
|
|
|
def get_update_df(df, merge_dict, mean_list, cls_list):
|
|
update_df = df.copy()
|
|
update_df["label"] = update_df["label"].apply(lambda x: x if x not in merge_dict.keys() else merge_dict[x])
|
|
color_dict = get_color_dict(mean_list, cls_list)
|
|
update_df["r"] = update_df.apply(lambda x: color_dict[x["label"]]["r"], axis=1)
|
|
update_df["g"] = update_df.apply(lambda x: color_dict[x["label"]]["g"], axis=1)
|
|
update_df["b"] = update_df.apply(lambda x: color_dict[x["label"]]["b"], axis=1)
|
|
return update_df, color_dict
|
|
|
|
def split_img_df(df, show=False):
|
|
img_list = []
|
|
for cls_no in tqdm(list(df["label"].unique())):
|
|
img_df = df.copy()
|
|
img_df.loc[df["label"] != cls_no, ["a"]] = 0
|
|
df_img = df2rgba(img_df)
|
|
if df_img is None:
|
|
continue
|
|
df_img = df_img.astype(np.uint8)
|
|
img_list.append(df_img)
|
|
return img_list
|
|
|
|
|
|
def get_base(img, loops, cls_num, threshold, size, h_split, v_split, n_cluster, alpha, th_rate, bg_split=True, debug=False):
|
|
if bg_split == False:
|
|
df = rgba2df(img)
|
|
df_list = [df]
|
|
else:
|
|
df_list = get_foreground(img, h_split, v_split, n_cluster, alpha, th_rate)
|
|
|
|
output_list = []
|
|
|
|
for idx, df in enumerate(df_list):
|
|
output_df = df.copy()
|
|
cls = MiniBatchKMeans(n_clusters = cls_num)
|
|
cls.fit(df[["r","g","b"]])
|
|
df["label"] = cls.labels_
|
|
df["label"] = df["label"].astype(str) + f"_{idx}"
|
|
for i in range(loops):
|
|
if i !=0:
|
|
img = df2rgba(df).astype(np.uint8)
|
|
blur_list, mean_list, cls_list = get_blur_cls(img, df["label"], size)
|
|
ciede_df = calc_ciede(mean_list, cls_list)
|
|
merge_dict = get_cls_update(ciede_df, df, threshold)
|
|
update_df, color_dict = get_update_df(df, merge_dict, mean_list, cls_list)
|
|
df = update_df
|
|
#if debug==True: img_plot(df)
|
|
output_df["label"] = df["label"]
|
|
output_df["layer_no"] = idx
|
|
output_list.append(output_df)
|
|
|
|
output_df = pd.concat(output_list).sort_index()
|
|
|
|
mean_list = []
|
|
cls_list = list(output_df["label"].unique())
|
|
for cls_no in tqdm(cls_list):
|
|
mask = get_mask(output_df, cls_no)
|
|
img_df = output_df.copy()
|
|
img_df.loc[output_df["label"] != cls_no, ["a"]] = 0
|
|
img_df, mean = fill_mean_color(img_df, mask)
|
|
mean_list.append(mean)
|
|
|
|
color_dict = get_color_dict(mean_list, cls_list)
|
|
output_df["r"] = output_df.apply(lambda x: color_dict[x["label"]]["r"], axis=1)
|
|
output_df["g"] = output_df.apply(lambda x: color_dict[x["label"]]["g"], axis=1)
|
|
output_df["b"] = output_df.apply(lambda x: color_dict[x["label"]]["b"], axis=1)
|
|
|
|
return output_df
|
|
|
|
def set_label(x, idx):
|
|
if x["m_flg"] == True:
|
|
return idx
|
|
else :
|
|
return x["label"]
|
|
|
|
def mode_fast(series):
|
|
return series.mode().iloc[0]
|
|
|
|
def get_seg_base(input_image, masks, th):
|
|
df = rgba2df(input_image)
|
|
df["label"] = -1
|
|
for idx, mask in tqdm(enumerate(masks)):
|
|
if int(mask["area"] < th):
|
|
continue
|
|
mask_df = mask2df(mask["segmentation"])
|
|
df = df.merge(mask_df, left_on=["x_l", "y_l"], right_on=["x_l_m", "y_l_m"], how="inner")
|
|
df["label"] = np.where(df["m_flg"] == True, idx, df["label"])
|
|
df.drop(columns=["x_l_m", "y_l_m", "m_flg"], inplace=True)
|
|
|
|
df['r'] = df.groupby('label')['r'].transform(mode_fast)
|
|
df['g'] = df.groupby('label')['g'].transform(mode_fast)
|
|
df['b'] = df.groupby('label')['b'].transform(mode_fast)
|
|
return df
|
|
|
|
def get_normal_layer(input_image, df):
|
|
base_layer_list = split_img_df(df, show=False)
|
|
|
|
org_df = rgba2df(input_image)
|
|
hsv_df = hsv2df(cv2.cvtColor(df2rgba(df).astype(np.uint8), cv2.COLOR_RGB2HSV))
|
|
hsv_org = hsv2df(cv2.cvtColor(input_image, cv2.COLOR_RGB2HSV))
|
|
|
|
hsv_org["bright_flg"] = hsv_df["v"] < hsv_org["v"]
|
|
bright_df = org_df.copy()
|
|
bright_df["bright_flg"] = hsv_org["bright_flg"]
|
|
bright_df["a"] = np.where(bright_df["bright_flg"] == True, 255, 0)
|
|
bright_df["label"] = df["label"]
|
|
bright_layer_list = split_img_df(bright_df, show=False)
|
|
|
|
hsv_org["shadow_flg"] = hsv_df["v"] >= hsv_org["v"]
|
|
shadow_df = rgba2df(input_image)
|
|
shadow_df["shadow_flg"] = hsv_org["shadow_flg"]
|
|
shadow_df["a"] = np.where(shadow_df["shadow_flg"] == True, 255, 0)
|
|
shadow_df["label"] = df["label"]
|
|
shadow_layer_list = split_img_df(shadow_df, show=True)
|
|
|
|
return base_layer_list, bright_layer_list, shadow_layer_list
|
|
|
|
|
|
def get_composite_layer(input_image, df):
|
|
base_layer_list = split_img_df(df, show=False)
|
|
|
|
org_df = rgba2df(input_image)
|
|
|
|
org_df["r"] = org_df["r"].apply(lambda x:int(x))
|
|
org_df["g"] = org_df["g"].apply(lambda x:int(x))
|
|
org_df["b"] = org_df["b"].apply(lambda x:int(x))
|
|
|
|
org_df["diff_r"] = df["r"] - org_df["r"]
|
|
org_df["diff_g"] = df["g"] - org_df["g"]
|
|
org_df["diff_b"] = df["b"] - org_df["b"]
|
|
|
|
org_df["shadow_flg"] = org_df.apply(
|
|
lambda x: True if x["diff_r"] >= 0 and x["diff_g"] >= 0 and x["diff_b"] >= 0 else False,
|
|
axis=1
|
|
)
|
|
org_df["screen_flg"] = org_df.apply(
|
|
lambda x: True if x["diff_r"] < 0 and x["diff_g"] < 0 and x["diff_b"] < 0 else False,
|
|
axis=1
|
|
)
|
|
|
|
|
|
shadow_df = org_df.copy()
|
|
shadow_df["a"] = org_df.apply(lambda x: 255 if x["shadow_flg"] == True else 0, axis=1)
|
|
|
|
shadow_df["r"] = shadow_df["r"].apply(lambda x: x*255)
|
|
shadow_df["g"] = shadow_df["g"].apply(lambda x: x*255)
|
|
shadow_df["b"] = shadow_df["b"].apply(lambda x: x*255)
|
|
|
|
shadow_df["r"] = (shadow_df["r"])/df["r"]
|
|
shadow_df["g"] = (shadow_df["g"])/df["g"]
|
|
shadow_df["b"] = (shadow_df["b"])/df["b"]
|
|
|
|
shadow_df["label"] = df["label"]
|
|
shadow_layer_list = split_img_df(shadow_df, show=True)
|
|
|
|
screen_df = org_df.copy()
|
|
|
|
screen_df["a"] = screen_df["screen_flg"].apply(lambda x: 255 if x == True else 0)
|
|
|
|
screen_df["r"] = (screen_df["r"] - df["r"])/(1 - df["r"]/255)
|
|
screen_df["g"] = (screen_df["g"] - df["g"])/(1 - df["g"]/255)
|
|
screen_df["b"] = (screen_df["b"] - df["b"])/(1 - df["b"]/255)
|
|
|
|
screen_df["label"] = df["label"]
|
|
screen_layer_list = split_img_df(screen_df, show=True)
|
|
|
|
|
|
addition_df = org_df.copy()
|
|
addition_df["a"] = addition_df.apply(lambda x: 255 if x["screen_flg"] == False and x["shadow_flg"] == False else 0, axis=1)
|
|
|
|
addition_df["r"] = org_df["r"] - df["r"]
|
|
addition_df["g"] = org_df["g"] - df["g"]
|
|
addition_df["b"] = org_df["b"] - df["b"]
|
|
|
|
addition_df["r"] = addition_df["r"].apply(lambda x: 0 if x < 0 else x)
|
|
addition_df["g"] = addition_df["g"].apply(lambda x: 0 if x < 0 else x)
|
|
addition_df["b"] = addition_df["b"].apply(lambda x: 0 if x < 0 else x)
|
|
|
|
addition_df["label"] = df["label"]
|
|
|
|
addition_layer_list = split_img_df(addition_df, show=True)
|
|
|
|
subtract_df = org_df.copy()
|
|
subtract_df["a"] = subtract_df.apply(lambda x: 255 if x["screen_flg"] == False and x["shadow_flg"] == False else 0, axis=1)
|
|
|
|
subtract_df["r"] = df["r"] - org_df["r"]
|
|
subtract_df["g"] = df["g"] - org_df["g"]
|
|
subtract_df["b"] = df["b"] - org_df["b"]
|
|
|
|
subtract_df["r"] = subtract_df["r"].apply(lambda x: 0 if x < 0 else x)
|
|
subtract_df["g"] = subtract_df["g"].apply(lambda x: 0 if x < 0 else x)
|
|
subtract_df["b"] = subtract_df["b"].apply(lambda x: 0 if x < 0 else x)
|
|
|
|
subtract_df["label"] = df["label"]
|
|
|
|
subtract_layer_list = split_img_df(subtract_df, show=True)
|
|
|
|
|
|
return base_layer_list, shadow_layer_list, screen_layer_list, addition_layer_list, subtract_layer_list
|