Spaces:
Build error
Build error
| # 功能函数 | |
| import os | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image, ImageFilter | |
| from matplotlib import pyplot as plt | |
| import collections | |
| import colorsys | |
| #定义字典存放颜色分量上下限 | |
| #例如：{颜色: [min分量, max分量]} | |
| #{'red': [array([160, 43, 46]), array([179, 255, 255])]} | |
| def getColorList(): | |
| dict = collections.defaultdict(list) | |
| # 黑色 | |
| lower_black = np.array([0, 0, 0]) | |
| upper_black = np.array([180, 255, 46]) | |
| color_list = [] | |
| color_list.append(lower_black) | |
| color_list.append(upper_black) | |
| dict['black'] = color_list | |
| #灰色 | |
| lower_gray = np.array([0, 0, 46]) | |
| upper_gray = np.array([180, 43, 220]) | |
| color_list = [] | |
| color_list.append(lower_gray) | |
| color_list.append(upper_gray) | |
| dict['gray']=color_list | |
| # 白色 | |
| lower_white = np.array([0, 0, 221]) | |
| upper_white = np.array([180, 30, 255]) | |
| color_list = [] | |
| color_list.append(lower_white) | |
| color_list.append(upper_white) | |
| dict['white'] = color_list | |
| # 红色1 | |
| lower_red = np.array([0, 43, 46]) | |
| upper_red = np.array([10, 255, 255]) | |
| color_list = [] | |
| color_list.append(lower_red) | |
| color_list.append(upper_red) | |
| dict['red'] = color_list | |
| #红色2 | |
| lower_red = np.array([156, 43, 46]) | |
| upper_red = np.array([180, 255, 255]) | |
| color_list = [] | |
| color_list.append(lower_red) | |
| color_list.append(upper_red) | |
| dict['red2']=color_list | |
| #橙色 | |
| lower_orange = np.array([11, 43, 46]) | |
| upper_orange = np.array([25, 255, 255]) | |
| color_list = [] | |
| color_list.append(lower_orange) | |
| color_list.append(upper_orange) | |
| dict['orange'] = color_list | |
| #黄色 | |
| lower_yellow = np.array([26, 43, 46]) | |
| upper_yellow = np.array([34, 255, 255]) | |
| color_list = [] | |
| color_list.append(lower_yellow) | |
| color_list.append(upper_yellow) | |
| dict['yellow'] = color_list | |
| #绿色 | |
| lower_green = np.array([35, 43, 46]) | |
| upper_green = np.array([77, 255, 255]) | |
| color_list = [] | |
| color_list.append(lower_green) | |
| color_list.append(upper_green) | |
| dict['green'] = color_list | |
| #青色 | |
| lower_cyan = np.array([78, 43, 46]) | |
| upper_cyan = np.array([99, 255, 255]) | |
| color_list = [] | |
| color_list.append(lower_cyan) | |
| color_list.append(upper_cyan) | |
| dict['cyan'] = color_list | |
| #蓝色 | |
| lower_blue = np.array([100, 43, 46]) | |
| upper_blue = np.array([124, 255, 255]) | |
| color_list = [] | |
| color_list.append(lower_blue) | |
| color_list.append(upper_blue) | |
| dict['blue'] = color_list | |
| # 紫色 | |
| lower_purple = np.array([125, 43, 46]) | |
| upper_purple = np.array([155, 255, 255]) | |
| color_list = [] | |
| color_list.append(lower_purple) | |
| color_list.append(upper_purple) | |
| dict['purple'] = color_list | |
| return dict | |
| # 输入hsv图片得到颜色 | |
| def get_color_from_pos(x, y, hsv, color_dict): | |
| assert x >= 0 and x < hsv.shape[1] | |
| assert y >= 0 and y < hsv.shape[0] | |
| point_hsv = hsv[int(y), int(x),:] | |
| h, s, v = point_hsv[0], point_hsv[1], point_hsv[2] | |
| for c in color_dict: | |
| low_rnk, high_rnk = color_dict[c][0], color_dict[c][1] | |
| if h >= low_rnk[0] and h <= high_rnk[0] and s >= low_rnk[1] and s <= high_rnk[1] and v >= low_rnk[2] and v <= high_rnk[2]: | |
| return c, point_hsv | |
| def get_color_from_image(hsv, color_dict): | |
| h, w, _ = hsv.shape | |
| center_hsv = hsv[int(h/2-1), int(w/2-1),:] | |
| h, s, v = center_hsv[0], center_hsv[1], center_hsv[2] | |
| print('source hsv', h, s, v) | |
| for c in color_dict: | |
| low_rnk, high_rnk = color_dict[c][0], color_dict[c][1] | |
| if h >= low_rnk[0] and h <= high_rnk[0] and s >= low_rnk[1] and s <= high_rnk[1] and v >= low_rnk[2] and v <= high_rnk[2]: | |
| return c, center_hsv | |
| def get_color_from_image_2(hsv, color_dict): | |
| hist, bins = np.histogram(hsv[:, :, 0], bins=180, range=(0,180)) | |
| hist[0] = 0 | |
| for i in range(len(hist)): | |
| if hist[i] == hist.max(): | |
| idx = i | |
| break | |
| print(idx) | |
| for c in color_dict: | |
| low_rnk, high_rnk = color_dict[c][0], color_dict[c][1] | |
| if idx >= low_rnk[0] and idx <= high_rnk[0]: | |
| return c, idx | |
| # 输入rgb值得到hsv值 | |
| def rgb2hsv(rgb, color_dict): | |
| assert len(rgb) == 3 | |
| r, g, b = rgb[0], rgb[1], rgb[2] | |
| hsv = colorsys.rgb_to_hsv(r/255, g/255, b/255) | |
| hsv = np.array(hsv) * np.array([180, 255, 255]) | |
| h, s, v = int(hsv[0]), int(hsv[1]), int(hsv[2]) | |
| print('target hsv', h, s, v) | |
| for c in color_dict: | |
| low_rnk, high_rnk = color_dict[c][0], color_dict[c][1] | |
| if h >= low_rnk[0] and h <= high_rnk[0] and s >= low_rnk[1] and s <= high_rnk[1] and v >= low_rnk[2] and v <= high_rnk[2]: | |
| return c, hsv | |
| # 判断颜色是否为彩色，黑白灰返回false | |
| def judge_color(color): | |
| if 'black' in color or 'gray' in color or 'white' in color: | |
| return False | |
| else: | |
| return True | |
| # 获得遮罩 | |
| def get_mask(color, hsv, color_dict): | |
| if 'red' in color: | |
| mask1 = cv2.inRange(hsv, color_dict['red'][0], color_dict['red'][1]) | |
| mask2 = cv2.inRange(hsv, color_dict['red2'][0], color_dict['red2'][1]) | |
| mask = mask1 + mask2 | |
| else: | |
| mask = cv2.inRange(hsv, color_dict[color][0], color_dict[color][1]) | |
| return mask != 0 | |
| # 保存mask | |
| def save_mask(path, mask, cnt): | |
| mask[mask!=0] = 255 | |
| mask = Image.fromarray(mask).convert('1') | |
| mask_name = os.path.join(path, str(cnt)+'.png') | |
| mask.save(mask_name) | |
| # 获得局部区域的属于color范围的遮罩 | |
| def get_region_mask(color, hsv, x1, y1, x, y, color_dict): | |
| if 'red' in color: | |
| mask1 = cv2.inRange(hsv, color_dict['red'][0], color_dict['red'][1]) | |
| mask2 = cv2.inRange(hsv, color_dict['red2'][0], color_dict['red2'][1]) | |
| mask = mask1 + mask2 | |
| else: | |
| mask = cv2.inRange(hsv, color_dict[color][0], color_dict[color][1]) | |
| mask = mask != 0 | |
| tmp = np.zeros(mask.shape) | |
| tmp[y1:y, x1:x] = mask[y1:y, x1:x] | |
| return tmp != 0 | |
| def get_region_mask_rect(color, hsv, x1, y1, x, y, color_dict): | |
| mask = cv2.inRange(hsv, color_dict[color][0], color_dict[color][1]) | |
| mask = mask == 0 # 除了这个颜色，其他的都置为True | |
| tmp = np.zeros(mask.shape) | |
| tmp[y1:y, x1:x] = mask[y1:y, x1:x] | |
| return tmp != 0 | |
| def get_region_mask_circle(color, hsv, black_back, color_dict): | |
| mask = cv2.inRange(hsv, color_dict[color][0], color_dict[color][1]) | |
| mask = mask == 0 # 除了这个颜色，其他的都置为True | |
| black_back = black_back!=0 | |
| mask = mask & black_back # 不是这个颜色且在这个区域 | |
| return mask | |
| def get_region_mask_remove(color, hsv, white_back, color_dict): # 去除涂抹区域中的目标颜色 | |
| white_back = white_back != 0 # 涂抹的区域 | |
| return white_back | |
| # 加光照函数 输入rgb 输出rgb | |
| def adding_light(img, light_conf): | |
| rows, cols = img.shape[:2] | |
| cX_list = [0, 0, 0, rows-1, rows-1, rows-1] | |
| cY_list = [0, cols // 2, cols - 1, 0, cols // 2, cols - 1] | |
| radius = min(rows, cols) | |
| strength = light_conf['light_strength'] | |
| cX = cX_list[light_conf['light_type']] | |
| cY = cY_list[light_conf['light_type']] | |
| dist = np.zeros((rows, cols)) | |
| x = np.arange(0, rows) | |
| dx = np.power((x - cX), 2) | |
| for i in range(cols): | |
| y = np.zeros(rows) + i | |
| dy = np.power((y - cY), 2) | |
| dist[:, i] = dx + dy | |
| R, G, B = img[:, :, 0], img[:, :, 1], img[:, :, 2] | |
| res = strength * (1.0 - np.sqrt(dist) / radius) | |
| mape = dist < radius * radius | |
| R = (R + res) * mape + R * (1 - mape) | |
| G = (G + res) * mape + G * (1 - mape) | |
| B = (B + res) * mape + B * (1 - mape) | |
| R = np.clip(R, 0, 255) | |
| G = np.clip(G, 0, 255) | |
| B = np.clip(B, 0, 255) | |
| res = np.uint8(np.dstack((R, G, B))) | |
| return res | |
| def Sort_map(result): | |
| idx = 0 | |
| res = {} | |
| for key in result.keys(): | |
| if "best" in key: | |
| res[str(idx)+'.png'] = result[key] | |
| idx += 1 | |
| for key in result.keys(): | |
| if "best" not in key: | |
| res[str(idx)+'.png'] = result[key] | |
| idx += 1 | |
| return res | |
| def get_name(flag_h, flag_s, flag_v): | |
| save_name = "" | |
| if flag_h != None: | |
| save_name += "h_" | |
| save_name += str(int(flag_h[0])) | |
| save_name += "_" | |
| save_name += str(int(flag_h[1])) | |
| save_name += "_" | |
| if flag_s != None: | |
| save_name += "s_" | |
| save_name += str(int(flag_s[0])) | |
| save_name += "_" | |
| save_name += str(int(flag_s[1])) | |
| save_name += "_" | |
| if flag_v != None: | |
| save_name += "v_" | |
| save_name += str(int(flag_v[0])) | |
| save_name += "_" | |
| save_name += str(int(flag_v[1])) | |
| save_name = save_name.strip('_') | |
| return save_name + '.jpg' | |
| def light_map(x, l_rnk, h_rnk): | |
| if x <= l_rnk: | |
| res = 0 | |
| elif x >= h_rnk: | |
| res = 255 | |
| else: | |
| k = 255 / (h_rnk - l_rnk) | |
| b = -255 * l_rnk / (h_rnk - l_rnk) | |
| res = k * x + b | |
| return int(res) | |
| def curve(hsv, mask, h_rnk, l_rnk): | |
| mask = np.array(mask) | |
| mask = mask > 10 | |
| v = hsv[:, :, 2] | |
| v_ = v.copy() | |
| vfunc = np.vectorize(light_map, excluded=['l_rnk', 'h_rnk']) | |
| v_ = vfunc(v_, l_rnk=l_rnk, h_rnk=h_rnk) | |
| hsv[:, :, 2] = v_ * mask + v * (~mask) | |
| return hsv | |
| # 保存结果 | |
| def save_result(hsv, alpha, mask_ori, pic_path, light_conf, curve_conf): | |
| if curve_conf['if_curve'] == True: | |
| hsv = curve(hsv=hsv, mask=mask_ori, h_rnk=curve_conf['h_rnk'], l_rnk=curve_conf['l_rnk']) | |
| if alpha == None: # 三通道图片 | |
| tag_img = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB) | |
| # | |
| # tag_img= cv2.GaussianBlur(tag_img, (9, 9), 0) | |
| if light_conf['if_light'] != False: | |
| tag_img = adding_light(tag_img, light_conf) | |
| res = Image.open(pic_path) | |
| res.paste(Image.fromarray(tag_img), mask = mask_ori) | |
| else: | |
| tag_img = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB) | |
| if light_conf['if_light'] != False: | |
| tag_img = adding_light(tag_img, light_conf) | |
| if mask_ori!= None: | |
| res = Image.open(pic_path) | |
| res.paste(Image.fromarray(tag_img), mask=mask_ori) | |
| r, g, b, _ = res.split() | |
| res = Image.merge('RGBA', (r, g, b, alpha)) | |
| else: | |
| r, g, b = Image.fromarray(tag_img[:, :, 0]), Image.fromarray(tag_img[:, :, 1]), Image.fromarray(tag_img[:, :, 2]) | |
| res = Image.merge('RGBA', (r, g, b, alpha)) | |
| # res.save(save_name) | |
| return res | |
| def mapping(mape, a, l_flag, h_flag, pre_set): | |
| if pre_set != None: | |
| assert len(pre_set) == 2 | |
| l_flag = pre_set[0] | |
| h_flag = pre_set[1] | |
| mi = mape[a].min() | |
| mx = mape[a].max() | |
| # 相交 | |
| if (l_flag < mx and l_flag > mi) or (h_flag > mi and h_flag < mx): | |
| if mx > h_flag: | |
| if mi >= l_flag: # [mi, h_flag] | |
| mape = (mape - mi) / (mx - mi) * (h_flag - mi) + mi | |
| else: # [l_flag, h_flag] | |
| mape = (mape - mi) / (mx - mi) * (h_flag - l_flag) + l_flag | |
| else: | |
| if mi < l_flag: # [l_flag, mx] | |
| mape = (mape - mi) / (mx - mi) * (mx - l_flag) + l_flag | |
| # 不相交 | |
| else: | |
| delta = l_flag - mi | |
| mape = mape + delta | |
| mi = mape[a].min() | |
| mx = mape[a].max() | |
| mape = (mape - mi) / (mx - mi) * (h_flag - l_flag) + l_flag | |
| return mape | |
| def mapping_diff(mape, mask_yes, mask_no, a): | |
| mask1 = a & mask_no # 车漆颜色 并且 不属于这个颜色 | |
| mask2 = a & mask_yes # 车漆颜色 并且 属于这个颜色 | |
| return | |
| def get_median(mape, mask): | |
| med = np.median(mape[mask]) | |
| mask_1_4 = (mape < med) & mask | |
| mask_3_4 = (mape > med) & mask | |
| med_1_4 = np.median(mape[mask_1_4]) | |
| med_3_4 = np.median(mape[mask_3_4]) | |
| return med_1_4, med_3_4, med | |
| # 汽车转色 | |
| def convert_color(hsv, mask, src_hsv, tag_hsv, tag_color, alpha, color_dict, mask_ori, pic_path, flag_s, flag_v, flag_h, light_conf, v_contrast, curve_conf): | |
| h, s, v = hsv[:, :, 0], hsv[:, :, 1], hsv[:, :, 2] | |
| tag_h, tag_s, tag_v = int(tag_hsv[0]), int(tag_hsv[1]), int(tag_hsv[2]) | |
| a = mask | |
| b = mask == 0 | |
| if 'white' in tag_color and flag_s == None: | |
| flag_s = [0, min(10, tag_s+1)] | |
| clahe = cv2.createCLAHE(clipLimit=v_contrast, tileGridSize=(8,8)) | |
| delta_h = tag_hsv[0] - src_hsv[0] | |
| delta_s = tag_hsv[1] - src_hsv[1] | |
| delta_v = tag_hsv[2] - src_hsv[2] | |
| print(f'delta_h: {delta_h}, delta_s: {delta_s}, delta_v: {delta_v}') | |
| h_ = h + delta_h | |
| s_ = s + delta_s | |
| v_ = v + delta_v | |
| print('flag_s:{}, flag_v:{}'.format(flag_s, flag_v)) | |
| # 映射到指定的区间 | |
| h_ = mapping(h_, a, color_dict[tag_color][0][0], color_dict[tag_color][1][0], flag_h) | |
| s_ = mapping(s_, a, color_dict[tag_color][0][1], color_dict[tag_color][1][1], flag_s) | |
| v_ = mapping(v_, a, color_dict[tag_color][0][2], color_dict[tag_color][1][2], flag_v) | |
| # 变彩色 色调选填 | |
| if judge_color(tag_color): # 只调色调, s, v通过设置参数来调整 | |
| if v_contrast != 0: | |
| v_1= clahe.apply((v_ * a).astype(np.uint8)) | |
| hsv[:, :, 2] = v_1 * a + v * b | |
| else: | |
| hsv[:, :, 2] = v_ * a + v * b | |
| hsv[:, :, 1] = s_ * a + s * b | |
| hsv[:, :, 0] = h_ * a + h * b | |
| res = save_result(hsv, alpha, mask_ori, pic_path, light_conf, curve_conf) | |
| # 变黑白灰 | |
| else: | |
| # 黑白灰色调[0,180]不用管 | |
| hsv[:, :, 0] = h_ * a + h * b | |
| hsv[:, :, 1] = s_ * a + s * b | |
| if v_contrast != 0: | |
| v_= clahe.apply((v_ * a).astype(np.uint8)) | |
| hsv[:, :, 2] = v_ * a + v * b | |
| res = save_result(hsv, alpha, mask_ori, pic_path, light_conf, curve_conf) | |
| return res | |