from CDM.detect_merge.Element import Element import CDM.detect_compo.lib_ip.ip_preprocessing as pre import time import cv2 import torch import numpy as np from torchvision import models from torch import nn import pandas as pd import re import openai import random import os from CDM.detect_merge.merge import reassign_ids import CDM.detect_merge.merge as merge from os.path import join as pjoin, exists label_dic ={'72':'Location', '42':'Photos', '77':'Social media', '91':'Voices', '6':'Email', '89':'Social media', '40':'Location', '43':'Phone', '82':'Photos', '3':'Contacts', '68':'Contacts', '49':'Profile', '56':'Photos'} keyword_list = {'Name':['name', 'first name', 'last name', 'full name', 'real name', 'surname', 'family name', 'given name'], 'Birthday':['birthday', 'date of birth', 'birth date', 'DOB', 'dob full birthday', 'birth year'], 'Address':['mailing address', 'physical address', 'postal address', 'billing address', 'shipping address', 'delivery address', 'residence', 'collect address', 'personal address', 'residential address'], 'Phone':['phone', 'phone number', 'mobile', 'mobile phone', 'mobile number', 'telephone', 'telephone number', 'call'], 'Email':['email', 'e-mail', 'email address', 'e-mail address'], 'Contacts':['contacts', 'phone-book', 'phone book', 'phonebook', 'contact list', 'phone contacts', 'address book'], 'Location':['location', 'locate', 'geography', 'geo', 'geo-location', 'precision location', 'nearby'], 'Photos':['camera', 'photo', 'scan', 'album', 'picture', 'gallery', 'photo library', 'storage', 'image', 'video', 'scanner', 'photograph'], 'Voices':['microphone', 'voice', 'mic', 'speech', 'talk'], 'Financial info':['credit card', 'pay', 'payment', 'debit card', 'mastercard', 'wallet'], 'IP':['IP', 'Internet Protocol', 'IP address', 'internet protocol address'], 'Cookies':['cookies', 'cookie'], 'Social media':['facebook', 'twitter', 'socialmedia', 'social media'], 'Profile':['profile', 'account'], 'Gender':['gender']} def summarize_segment(segment, label, word_limit=20): openai.api_key = os.environ.get('openai_key') # 在提示中明确要求缩略后的段落字数不超过给定的字数 prompt = ( f"Extract key information from this paragraph related to the topic '{label}' " f"in no more than {word_limit} words: \"{str(segment)}\". Without using phrases like 'Key information'. " ) response = openai.ChatCompletion.create( model="gpt-4o-mini", messages=[ {"role": "user", "content": prompt} ], max_tokens=400, # 可以根据需要调整，确保能生成适当的长度 n=1, stop=None, temperature=0, ) shortened_segment = response.choices[0].message['content'].strip() return shortened_segment def get_data_type(sentence, keywords, use_gpt=True): sent_data_type = "others" if use_gpt: openai.api_key = os.environ["OPENAI_API_KEY"] prompt = f"Is this piece of texts \"{sentence}\" related to any following privacy information data types? Or not relevant to any of them? ONLY answer the data type or \"not relevant\". ONLY use following data type list. Data types and their Description:\n" \ f"Name: How a user refers to themselves," \ f" Birthday: A user’s birthday," \ f" Address: A user’s address," \ f" Phone: A user’s phone number," \ f" Email: A user’s email address," \ f" Contacts: A user’s contact information, or the access to the contact permission," \ f" Location: A user’s location information, or the access to the location permission," \ f" Photos: A user’s photos, videos, or the access to the camera permission," \ f" Voices: A user’s voices, recordings, or the access to the microphone permission," \ f" Financial Info: Information about a user’s financial accounts, purchases, or transactions," \ f" Profile: A user’s account information," \ f"Social Media: A user's social media information, or the access to social media accounts" response = openai.ChatCompletion.create( # engine="text-davinci-002", model="gpt-4o-mini", messages=[ # {"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": prompt} ], max_tokens=100, n=1, stop=None, temperature=0, ) # response_full_text = response.choices[0].text.strip() response_full_text = response.choices[0].message['content'] for k in keywords.keys(): if k == "Financial info" or k == "Social media": if k.lower() in response_full_text.lower(): sent_data_type = k break else: words = re.split(r'\W+', response_full_text.lower()) if k.lower() in words: sent_data_type = k break # print("----------------------") # print("sentence: ", sentence) # print("prompt: ", prompt) # print("response: ", response_full_text) # print("sent_data_type: ", sent_data_type) else: for k in keywords.keys(): for w in keywords[k]: words = re.split(r'\W+', sentence.lower()) if w.lower() in words: sent_data_type = k break if sent_data_type != "others": break return sent_data_type # def get_clf_model(use_resnet18=True, use_gpu=False): # # device = 'cpu' # if use_gpu: # device = 'cuda:0' # # if use_resnet18: # model = models.resnet18().to(device) # in_feature_num = model.fc.in_features # model.fc = nn.Linear(in_feature_num, 99) # model.conv1 = nn.Conv2d(in_channels=1, out_channels=64, kernel_size=(5, 5), padding=(3, 3), stride=(2, 2), # bias=False) # # PATH = "./CDM/model/model-99-resnet18.pkl" # model.load_state_dict(torch.load(PATH, map_location=torch.device(device))) # # model.eval() # else: # # replace with your own model # None # # return model def get_clf_model(clf_model="ResNet18", use_gpu=False): device = 'cpu' if use_gpu: device = 'cuda:0' if clf_model == "ResNet18": model = models.resnet18().to(device) in_feature_num = model.fc.in_features model.fc = nn.Linear(in_feature_num, 99) model.conv1 = nn.Conv2d(in_channels=1, out_channels=64, kernel_size=(5, 5), padding=(3, 3), stride=(2, 2), bias=False) PATH = "./CDM/model/model-99-resnet18.pkl" model.load_state_dict(torch.load(PATH, map_location=torch.device(device))) model.eval() elif clf_model == "ViT": model = torch.load('./CDM/model/model-99-ViT-entire.pkl', map_location=torch.device(device)) model = model.to(device) model.eval() else: # replace with your own model None return model def compo_classification(input_img, output_root, segment_root, merge_json, output_data, resize_by_height=800, clf_model="ResNet18", model=get_clf_model("ResNet18")): # load text and non-text compo ele_id = 0 compos = [] texts = [] elements = [] for compo in merge_json['compos']: if compo['class'] == 'Text': element = Element(ele_id, (compo["position"]['column_min'], compo["position"]['row_min'], compo["position"]['column_max'], compo["position"]['row_max']), 'Text', text_content=compo['text_content']) texts.append(element) ele_id += 1 else: element = Element(ele_id, (compo["position"]['column_min'], compo["position"]['row_min'], compo["position"]['column_max'], compo["position"]['row_max']), compo['class']) compos.append(element) ele_id += 1 org, grey = pre.read_img(input_img, resize_by_height) grey = grey.astype('float32') grey = grey / 255 # grey = (grey - grey.mean()) / grey.std() # --------- classification ---------- # for compo in compos: # # # comp_grey = grey[compo.row_min:compo.row_max, compo.col_min:compo.col_max] # # # # comp_crop = cv2.resize(comp_grey, (32, 32)) # # # # comp_crop = comp_crop.reshape(1, 1, 32, 32) # # # # comp_tensor = torch.tensor(comp_crop) # # comp_tensor = comp_tensor.permute(0, 1, 3, 2) # # # # model = get_clf_model() # # pred_label = model(comp_tensor) # # # # if str(np.argmax(pred_label.cpu().data.numpy(), axis=1)[0]) in label_dic.keys(): # # compo.label = label_dic[str(np.argmax(pred_label.cpu().data.numpy(), axis=1)[0])] # # elements.append(compo) # # else: # # compo.label = str(np.argmax(pred_label.cpu().data.numpy(), axis=1)[0]) # # if clf_model == "ResNet18": # # comp_grey = grey[compo.row_min:compo.row_max, compo.col_min:compo.col_max] # # comp_crop = cv2.resize(comp_grey, (32, 32)) # # comp_crop = comp_crop.reshape(1, 1, 32, 32) # # comp_tensor = torch.tensor(comp_crop) # comp_tensor = comp_tensor.permute(0, 1, 3, 2) # # # model = get_clf_model(clf_model) # pred_label = model(comp_tensor) # # if str(np.argmax(pred_label.cpu().data.numpy(), axis=1)[0]) in label_dic.keys(): # compo.label = label_dic[str(np.argmax(pred_label.cpu().data.numpy(), axis=1)[0])] # elements.append(compo) # else: # compo.label = str(np.argmax(pred_label.cpu().data.numpy(), axis=1)[0]) # # elif clf_model == "ViT": # # comp_grey = grey[compo.row_min:compo.row_max, compo.col_min:compo.col_max] # # comp_crop = cv2.resize(comp_grey, (224, 224)) # # # Convert the image to tensor # comp_tensor = torch.from_numpy(comp_crop) # # # Reshape and repeat along the channel dimension to convert to RGB # comp_tensor = comp_tensor.view(1, 224, 224).repeat(3, 1, 1) # # # comp_tensor = comp_tensor.permute(0, 2, 1) # # comp_tensor = comp_tensor.unsqueeze(0) # add a batch dimension # # # model = get_clf_model(clf_model) # # pred_label = model(comp_tensor) # # # Forward pass through the model # with torch.no_grad(): # output = model(comp_tensor) # # # Get the predicted label # _, predicted = torch.max(output.logits, 1) # # # print("predicted_label: ", predicted.cpu().numpy()) # # if str(predicted.cpu().numpy()[0]) in label_dic.keys(): # compo.label = label_dic[str(predicted.cpu().numpy()[0])] # elements.append(compo) # else: # compo.label = str(predicted.cpu().numpy()[0]) # # else: # print("clf_model has to be ResNet18 or ViT") # ============================================================================= # classification_start_time = time.time() # model = get_clf_model(clf_model) # # for compo in compos: # compo_start_time = time.time() # # # 计时预处理部分 # preprocess_start = time.time() # comp_grey = grey[compo.row_min:compo.row_max, compo.col_min:compo.col_max] # preprocess_time = time.time() - preprocess_start # # # 计时图像调整部分 # resize_start = time.time() # if clf_model == "ResNet18": # comp_crop = cv2.resize(comp_grey, (32, 32)) # elif clf_model == "ViT": # comp_crop = cv2.resize(comp_grey, (224, 224)) # resize_time = time.time() - resize_start # # # 计时张量转换部分 # tensor_start = time.time() # if clf_model == "ResNet18": # comp_crop = comp_crop.reshape(1, 1, 32, 32) # comp_tensor = torch.tensor(comp_crop) # comp_tensor = comp_tensor.permute(0, 1, 3, 2) # elif clf_model == "ViT": # comp_tensor = torch.from_numpy(comp_crop) # comp_tensor = comp_tensor.view(1, 224, 224).repeat(3, 1, 1) # comp_tensor = comp_tensor.unsqueeze(0) # tensor_time = time.time() - tensor_start # # # 计时模型推理部分 # inference_start = time.time() # with torch.no_grad(): # if clf_model == "ResNet18": # pred_label = model(comp_tensor) # output = pred_label # elif clf_model == "ViT": # output = model(comp_tensor) # inference_time = time.time() - inference_start # # # 计时后处理部分 # postprocess_start = time.time() # if clf_model == "ResNet18": # predicted = np.argmax(output.cpu().data.numpy(), axis=1)[0] # elif clf_model == "ViT": # _, predicted = torch.max(output.logits, 1) # predicted = predicted.cpu().numpy()[0] # # if str(predicted) in label_dic.keys(): # compo.label = label_dic[str(predicted)] # elements.append(compo) # else: # compo.label = str(predicted) # postprocess_time = time.time() - postprocess_start # # compo_total_time = time.time() - compo_start_time # # # 输出每个部分的耗时 # print("==============================================") # print(f"Component processing time: {compo_total_time:.4f}s") # print(f" Preprocessing time: {preprocess_time:.4f}s") # print(f" Resize time: {resize_time:.4f}s") # print(f" Tensor conversion time: {tensor_time:.4f}s") # print(f" Inference time: {inference_time:.4f}s") # print(f" Post-processing time: {postprocess_time:.4f}s\n") # print("==============================================") # ============================================================================= # model = get_clf_model(clf_model) classification_start_time = time.time() comp_tensors = [] elements = [] # 收集所有张量 for compo in compos: # 预处理 comp_grey = grey[compo.row_min:compo.row_max, compo.col_min:compo.col_max] # 调整图像大小 if clf_model == "ResNet18": comp_crop = cv2.resize(comp_grey, (32, 32)) elif clf_model == "ViT": comp_crop = cv2.resize(comp_grey, (224, 224)) # 张量转换 if clf_model == "ResNet18": comp_crop = comp_crop.reshape(1, 1, 32, 32) comp_tensor = torch.tensor(comp_crop).permute(0, 1, 3, 2) elif clf_model == "ViT": comp_tensor = torch.from_numpy(comp_crop) comp_tensor = comp_tensor.view(1, 224, 224).repeat(3, 1, 1).unsqueeze(0) comp_tensors.append(comp_tensor) # 将张量堆叠成批次 batch_tensor = torch.cat(comp_tensors, dim=0) # 模型推理 with torch.no_grad(): if clf_model == "ResNet18": output = model(batch_tensor) elif clf_model == "ViT": output = model(batch_tensor) # 后处理 if clf_model == "ResNet18": predicted = np.argmax(output.cpu().numpy(), axis=1) elif clf_model == "ViT": _, predicted = torch.max(output.logits, 1) predicted = predicted.cpu().numpy() # 为组件分配标签 for idx, compo in enumerate(compos): pred_label = predicted[idx] if str(pred_label) in label_dic.keys(): compo.label = label_dic[str(pred_label)] elements.append(compo) else: compo.label = str(pred_label) time_cost_ic = time.time() - classification_start_time print("time cost for icon classification: %2.2f s" % time_cost_ic) print("time cost for each icon classification: %2.2f s" % (time_cost_ic/len(compos))) # ic_time_cost_all.append(time_cost_ic) # --------- end classification ---------- text_selection_time = time.time() for this_text in texts: # found_flag = 0 # # for key in keyword_list: # for w in keyword_list[key]: # words = re.split(r'\W+', this_text.text_content.lower()) # if w.lower() in words: # this_text.label = key # elements.append(this_text) # found_flag = 1 # break # # if found_flag == 0: # this_text.label = 'others' retries = 10 for i in range(retries): try: text_label = get_data_type(this_text.text_content.lower(), keyword_list, use_gpt=False) break except openai.error.RateLimitError as e: if "overloaded" in str(e): # Exponential backoff with jitter sleep_time = 2 * (2 ** i) + random.uniform(0, 0.1) time.sleep(sleep_time) else: raise except Exception as e: raise this_text.label = text_label if this_text.label != "others": elements.append(this_text) time_cost_ts = time.time() - text_selection_time print("time cost for text selection: %2.2f s" % time_cost_ts) classification_cost = time.time() - classification_start_time # ts_time_cost_all.append(time_cost_ts) # ---------- end ------------------------------- # ---------- matching result ----------- index = input_img.split('/')[-1][:-4] app_id = str(index).split('-')[0] index_path = pjoin(segment_root, app_id, 'classified_sentences/keyword_index.txt') dict_index = {} if exists(index_path): with open(index_path, 'r') as g: for line in g: key, value = line.strip().split(':', 1) dict_index[key] = value for item in elements: complete_path = pjoin(segment_root, app_id, 'classified_sentences', item.label + '.txt') # print("complete_path: ", complete_path) if exists(complete_path): with open(complete_path, 'r', encoding='utf-8') as file: content = file.read() # Replace line breaks with spaces and strip any extra whitespace this_text = ' '.join(content.splitlines()).strip() lines = content.splitlines() non_empty_lines = [line for line in lines if line.strip() != ""] for i in range(len(non_empty_lines)): if non_empty_lines[i][0].isalpha(): non_empty_lines[i] = non_empty_lines[i][0].upper() + non_empty_lines[i][1:] # output_data = output_data.append({'screenshot': 's' + str(index), 'id': item.id + 1, 'label': item.label, 'index': dict_index[item.label], 'text': this_text, 'sentences': non_empty_lines}, ignore_index=True) output_data = pd.concat([output_data, pd.DataFrame([{'screenshot': 's' + str(index), 'id': item.id + 1, 'label': item.label, 'index': dict_index[item.label], 'text': this_text, 'sentences': non_empty_lines}])]) else: # output_data = output_data.append({'screenshot': 's' + str(index), 'id': item.id + 1, 'label': item.label, 'index': "None", 'text': "No information!", 'sentences': "None"}, # ignore_index=True) output_data = pd.concat([output_data, pd.DataFrame([{'screenshot': 's' + str(index), 'id': item.id + 1, 'label': item.label, 'index': "None", 'text': "No information!", 'sentences': "None"}])]) # -----------------openai缩略text--------------------------- short_text_time_start = time.time() short_output_data = output_data[['label', 'text']].copy() short_output_data = short_output_data.rename(columns={'text': 'segment'}) retries = 3 summarized_data = [] # List to hold summarized rows for index, row in short_output_data.iterrows(): seg = row['segment'] label = row['label'] if seg == "No information!": shortened_seg = seg else: for i in range(retries): try: shortened_seg = summarize_segment(seg, label) # print(seg) # print("--------------------") # print(shortened_seg) break except openai.error.RateLimitError as e: if "overloaded" in str(e): print("error") sleep_time = 2 * (2 ** i) + 0.1 # sleep_time = 3 time.sleep(sleep_time) except Exception as e: # If you wish, you can print or log the exception details here without raising it print(e) else: # This part will be executed if the for loop doesn't hit 'break' shortened_seg = seg summarized_data.append({'label': row['label'], 'segment': shortened_seg}) short_text_time_cost = time.time() - short_text_time_start print("缩短policy总共花费了： %2.2f s" % short_text_time_cost) # ------------------------------------------------------- full_size_org, full_size_grey = pre.read_img(input_img) ratio = full_size_org.shape[0]/org.shape[0] show = False wait_key = 0 reassign_ids(elements) board = merge.show_elements(full_size_org, elements, ratio, show=show, win_name='elements after merging', wait_key=wait_key, line=3) draw_pic_start = time.time() board_one_element = merge.show_one_element_class(full_size_org, elements, ratio, show=show, win_name='elements after merging', wait_key=wait_key, line=3, summarized_data=summarized_data) draw_pic_cost = time.time() - draw_pic_start print("生成所有的展示图片花费： %2.2f s" % draw_pic_cost) print("平均生成一张图片花费： %2.2f s" % (draw_pic_cost/len(board_one_element))) classification_root = pjoin(output_root, 'classification') # save all merged elements, clips and blank background name = input_img.replace('\\', '/').split('/')[-1][:-4] components = merge.save_elements(pjoin(classification_root, name + '.json'), elements, full_size_org.shape, ratio) cv2.imwrite(pjoin(classification_root, name + '.jpg'), board) print("一共图片张数： ", len(board_one_element)) for i in range(len(board_one_element)): e_name = str(int(elements[i].id) + 1) cv2.imwrite(pjoin(classification_root + '/GUI', name + '-' + e_name + '.jpg'), board_one_element[i]) print("生成图片展示总花费： %2.2f s" % (time.time()-short_text_time_start)) # print('[Classification Completed] Input: %s Output: %s' % (input_img, pjoin(classification_root, name + '.jpg'))) return time_cost_ic, time_cost_ts, output_data, board_one_element,classification_cost