# author : wangrui
# data : 2024/3/25 9:52
import time
from transformers import DetrImageProcessor, DetrForObjectDetection
import torch
from PIL import Image
import requests
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from sqlalchemy import create_engine, text
import pandas as pd
import concurrent.futures
import json
import threading
from googletrans import Translator
from requests.exceptions import RequestException
Image.MAX_IMAGE_PIXELS = 200000000 # 设置更高的限制
device = torch.device("cpu")
# 历史数据
class ImageRecognitionHistory(object):
def __init__(self):
self.mysql_host = "rm-wz9yg9bsb2zf01ea4yo.mysql.rds.aliyuncs.com"
self.mysql_port = "3306"
self.mysql_username = "adv_yswg"
self.mysql_pwd = "HCL1zcUgQesaaXNLbL37O5KhpSAy0c"
self.db_mame = "selection"
self.mysql_engine = create_engine(
f"mysql+pymysql://{self.mysql_username}:{self.mysql_pwd}@{self.mysql_host}:{self.mysql_port}/{self.db_mame}",
pool_size=30, # 连接池的大小
max_overflow=30 # 超出连接池大小之外可以创建的连接数
)
# self.processor = DetrImageProcessor.from_pretrained(r"E:\Model\model_new\detr-resnet-50", revision="no_timm")
# self.model = DetrForObjectDetection.from_pretrained(r"E:\Model\model_new\detr-resnet-50", revision="no_timm")
# self.image_classification = pipeline(Tasks.image_classification,
# model='E:\Model\model_new\cv_nextvit-small_image-classification_Dailylife-labels')
# 模型路径： hadoop15上 192.168.10.224
self.processor = DetrImageProcessor.from_pretrained(r"/home/wangrui/detr-resnet-50", revision="no_timm")
self.model = DetrForObjectDetection.from_pretrained(r"/home/wangrui/detr-resnet-50", revision="no_timm")
self.image_classification = pipeline(Tasks.image_classification,
model='/home/wangrui/cv_nextvit-small_image-classification_Dailylife-labels')
self.df_image = pd.DataFrame()
self.df_image_sub_list = []
self.threads_num = 25
def rea_and_prepare_data(self):
print("开始准备图片数据")
query = "SELECT id, files FROM file_files_test where module = 'Material' and english_elements is null"
self.df_image = pd.read_sql(query, self.mysql_engine)
sub_size = len(self.df_image) // self.threads_num
print("数据总量是：", len(self.df_image))
if self.threads_num == 1:
self.df_image_sub_list.append(self.df_image)
return
if sub_size:
for i in range(self.threads_num):
start_idx = i * sub_size
end_idx = (i + 1) * sub_size if i < (self.threads_num - 1) else len(self.df_image)
sub_df = self.df_image.iloc[start_idx: end_idx]
self.df_image_sub_list.append(sub_df)
print("数据已准备好")
else:
print("没有信息，请注意检查！")
quit(1)
def thread_task(self, sub_df):
print(f"线程-{threading.current_thread().ident} 开始处理数据")
conn = self.mysql_engine.connect()
count = 0
for index, row in sub_df.iterrows():
try:
parse_value = json.loads(row['files'])
img_path = parse_value['raw']
image_url = f"http://soundasia.oss-cn-shenzhen.aliyuncs.com/{img_path}"
english_elements, chinese_elements = self.get_image_element(image_url)
if english_elements and chinese_elements:
id = row['id']
# 构造参数绑定的 SQL 语句模板
sql_template = text(f"""
UPDATE file_files_test
SET english_elements = '{english_elements}',
chinese_elements = '{chinese_elements}'
WHERE id = {id}
""")
conn.execute(sql_template)
except Exception as e:
print(f"处理数据时出现异常: {e}")
finally:
count += 1
if count % 200 == 0 and count > 0:
print(f"线程-{threading.current_thread().ident} 已经处理了 {count}条数据")
print(f"线程-{threading.current_thread().ident} 处理数据完成")
def handle_image_history(self):
print("开始图片元素标注")
start_time = time.time()
with concurrent.futures.ThreadPoolExecutor(25) as executor:
futures = [executor.submit(self.thread_task, sub_df) for sub_df in self.df_image_sub_list]
for future in concurrent.futures.as_completed(futures):
future.result() # 等待每个线程任务完成
print("标签标注完成！")
elapsed_time = time.time() - start_time
print(f"耗时：{elapsed_time:.2f} 秒")
# def handle_image_history(self):
# print("开始图片元素标注")
# self.thread_task(self.df_image_sub_list[0])
# print("标签标注完成！")
def translate_with_retry(self, translator, text, src='en', dest='zh-CN', max_retries=10):
retries = 0
while retries < max_retries:
try:
translation = translator.translate(text, src=src, dest=dest)
return translation.text
except RequestException as e:
print(f"Request failed: {e}")
retries += 1
if retries < max_retries:
print(f"Retrying... ({retries}/{max_retries})")
time.sleep(2 ** retries) # Exponential backoff for retries
else:
print("Max retries reached. Translation failed.")
return None
def get_image_element(self, url):
try:
translator = Translator()
image = Image.open(requests.get(url, stream=True, timeout=30).raw)
if image.mode != "RGB":
image = image.convert("RGB")
inputs = self.processor(images=image, return_tensors="pt")
outputs = self.model(**inputs)
target_sizes = torch.tensor([image.size[::-1]])
results = self.processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.7)[0]
except Exception as e:
print(f"该图片解析失败: {e}")
return None, None
englist_label_set = set()
chinese_lable_set = set()
for label in results['labels']:
label = self.model.config.id2label[label.item()]
if label:
englist_label_set.add(label)
chinese_label = self.translate_with_retry(translator, label, src='en', dest='zh-CN')
if chinese_label:
chinese_lable_set.add(chinese_label)
if not englist_label_set:
results = self.image_classification(url)
for score, label in zip(results['scores'], results['labels']):
if score > 0.5:
english_label = self.translate_with_retry(translator, label, src='zh-CN', dest='en')
if english_label:
englist_label_set.add(english_label)
chinese_lable_set.add(label)
if englist_label_set:
return ','.join(englist_label_set), ','.join(chinese_lable_set)
else:
return None, None
def run(self):
self.rea_and_prepare_data()
self.handle_image_history()
if __name__ == '__main__':
handle_obj = ImageRecognitionHistory()
handle_obj.run()