299 lines
12 KiB
Python
299 lines
12 KiB
Python
import torch
|
|
from PIL import Image
|
|
from transformers import AutoModel, AutoTokenizer
|
|
import json
|
|
import re
|
|
from pymongo import MongoClient
|
|
import time
|
|
from bson import ObjectId
|
|
import os
|
|
import glob
|
|
from datetime import datetime, timedelta
|
|
|
|
# 数据库连接模块
|
|
class DatabaseHandler:
|
|
def __init__(self, mongo_uri, database_name, results_collection_name):
|
|
self.client = MongoClient(mongo_uri)
|
|
self.db = self.client[database_name]
|
|
self.results_collection = self.db[results_collection_name]
|
|
|
|
def save_result(self, result):
|
|
# 如果 result 中没有 filename，使用时间戳作为替代
|
|
# filename = result.get('filename', f"unknown_{result['timestamp']}")
|
|
filename = result.get('filename')
|
|
# 检查是否已存在相同 filename 的结果
|
|
existing_result = self.results_collection.find_one({'filename': filename})
|
|
if existing_result:
|
|
print(f"Video with filename {filename} has already been processed. Skipping.")
|
|
return
|
|
|
|
# 确保 result 中有 filename
|
|
result['filename'] = filename
|
|
|
|
# 将 ObjectId 转换为字符串
|
|
if 'video_id' in result and isinstance(result['video_id'], ObjectId):
|
|
result['video_id'] = str(result['video_id'])
|
|
|
|
self.results_collection.insert_one(result)
|
|
|
|
def is_sequence_processed(self, filename):
|
|
return self.results_collection.find_one({'filename': filename}) is not None
|
|
|
|
|
|
class JSONEncoder(json.JSONEncoder):
|
|
def default(self, o):
|
|
if isinstance(o, ObjectId):
|
|
return str(o)
|
|
return super().default(o)
|
|
|
|
# 视频处理模块
|
|
class ImageSequenceProcessor:
|
|
def __init__(self, model_dir):
|
|
self.model = AutoModel.from_pretrained(model_dir, trust_remote_code=True,
|
|
attn_implementation='sdpa', torch_dtype=torch.bfloat16).eval().cuda()
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
|
|
self.max_size = 512 # 设置最大尺寸
|
|
|
|
def compress_image(self, image):
|
|
# 保持纵横比的情况下调整图片大小
|
|
image.thumbnail((self.max_size, self.max_size))
|
|
|
|
# 如果图像已经是JPEG格式，直接返回调整大小后的图像
|
|
if image.format == 'JPEG':
|
|
return image
|
|
|
|
# 对于非JPEG格式，进行压缩
|
|
buffer = io.BytesIO()
|
|
if image.mode in ('RGBA', 'LA') or (image.mode == 'P' and 'transparency' in image.info):
|
|
# 保持透明度
|
|
image.save(buffer, format="PNG", optimize=True)
|
|
else:
|
|
# 转换为JPEG并压缩
|
|
image.convert('RGB').save(buffer, format="JPEG", quality=85, optimize=True)
|
|
|
|
buffer.seek(0)
|
|
return Image.open(buffer)
|
|
|
|
def process_image_sequence(self, image_paths):
|
|
frames = [self.compress_image(Image.open(img_path)) for img_path in image_paths]
|
|
question = "Analyze these 10 images as if they were frames from a video. Describe the scene in detail in Chinese, including the setting, number of people, their actions, and any changes or movements observed across the frames."
|
|
msgs = [
|
|
{'role': 'user', 'content': frames + [question]},
|
|
]
|
|
|
|
params = {
|
|
"use_image_id": False,
|
|
"max_slice_nums": 1
|
|
}
|
|
|
|
answer = self.model.chat(
|
|
image=None,
|
|
msgs=msgs,
|
|
tokenizer=self.tokenizer,
|
|
**params
|
|
)
|
|
|
|
extracted_info = self.extract_info(answer)
|
|
|
|
return {
|
|
"original_answer": answer,
|
|
"extracted_info": extracted_info,
|
|
"num_frames": len(frames),
|
|
}
|
|
|
|
@staticmethod
|
|
def extract_info(answer):
|
|
info = {
|
|
"environment": None,
|
|
"num_people": None,
|
|
"actions": [],
|
|
"interactions": [],
|
|
"objects": [],
|
|
"furniture": []
|
|
}
|
|
|
|
# 环境提取
|
|
environments = ["办公室", "室内", "室外", "会议室", "办公"]
|
|
for env in environments:
|
|
if env in answer.lower():
|
|
info["environment"] = env
|
|
break
|
|
|
|
# 改进的人数提取
|
|
people_patterns = [
|
|
r'(\d+)\s*(人|个人|位|名|员工|用户|小朋友|成年人|女性|男性)',
|
|
r'(一|二|三|四|五|六|七|八|九|十)\s*(人|个人|位|名|员工|用户|小朋友|成年人|女性|男性)',
|
|
r'(一个|几个)\s*(人|个人|员工|用户|小朋友|成年人|女性|男性)',
|
|
r'几\s*(名|位)\s*(人|员工|用户|小朋友|成年人|女性|男性)?',
|
|
r'(男|女)(性|生|士)',
|
|
r'(成年|未成年|青少年|老年)\s*(人|群体)',
|
|
r'(员工|职工|工人|学生|顾客|观众|游客|乘客)',
|
|
r'(群众|民众|大众|公众)',
|
|
r'(男女|老少|老幼|大人|小孩)'
|
|
]
|
|
for pattern in people_patterns:
|
|
match = re.search(pattern, answer)
|
|
if match:
|
|
if match.group(1).isdigit():
|
|
info["num_people"] = int(match.group(1))
|
|
elif match.group(1) in ['一个', '一']:
|
|
info["num_people"] = 1
|
|
else:
|
|
num_word_to_digit = {
|
|
'二': 2, '三': 3, '四': 4, '五': 5,
|
|
'六': 6, '七': 7, '八': 8, '九': 9, '十': 10
|
|
}
|
|
info["num_people"] = num_word_to_digit.get(match.group(1), 0)
|
|
break
|
|
|
|
# 动作和互动提取
|
|
actions = ["坐", "站", "摔倒", "跳舞", "转身", "摔", "倒", "倒下", "躺下", "转身", "跳跃", "跳", "躺", "睡", "说话"]
|
|
for action in actions:
|
|
if action in answer:
|
|
info["actions"].append(action)
|
|
|
|
interactions = ["互动", "交流", "身体语言", "交谈", "讨论", "开会"]
|
|
for interaction in interactions:
|
|
if interaction in answer:
|
|
info["interactions"].append(interaction)
|
|
|
|
# 物体和家具提取
|
|
objects = ["水瓶", "办公用品", "文件", "电脑"]
|
|
furniture = ["椅子", "桌子", "咖啡桌", "文件柜", "床", "沙发"]
|
|
|
|
for obj in objects:
|
|
if obj in answer:
|
|
info["objects"].append(obj)
|
|
|
|
for item in furniture:
|
|
if item in answer:
|
|
info["furniture"].append(item)
|
|
|
|
return info
|
|
|
|
# 主处理类
|
|
class ImageSequenceAnalysisSystem:
|
|
def __init__(self, mongo_uri, db_name, model_dir, results_collection_name):
|
|
self.db_handler = DatabaseHandler(mongo_uri, db_name, results_collection_name)
|
|
self.image_processor = ImageSequenceProcessor(model_dir)
|
|
self.last_processed_time = datetime.now() - timedelta(hours=1)
|
|
|
|
def get_all_images(self, image_folders):
|
|
image_files = []
|
|
for folder in image_folders:
|
|
image_files.extend(glob.glob(os.path.join(folder, '*.jpg')))
|
|
image_files.sort()
|
|
return image_files
|
|
|
|
|
|
def process_image_sequence(self, image_paths):
|
|
print(f"Attempting to process sequence: {[os.path.basename(img) for img in image_paths]}")
|
|
start_time = time.time()
|
|
try:
|
|
# 使用第一张图片的文件名作为序列的标识符
|
|
filename = os.path.basename(image_paths[0])
|
|
|
|
if self.db_handler.is_sequence_processed(filename):
|
|
print(f"Skipping already processed image sequence: {filename}")
|
|
return False
|
|
|
|
print("Processing new image sequence...")
|
|
result = self.image_processor.process_image_sequence(image_paths)
|
|
|
|
# timestamp = datetime.now()
|
|
# result['timestamp'] = timestamp.strftime("%Y%m%d_%H%M%S")
|
|
result['image_paths'] = image_paths
|
|
result['filename'] = filename
|
|
|
|
# 计算图片序列的周期
|
|
image_times = [self.get_file_time(img) for img in image_paths]
|
|
if len(image_times) >= 2:
|
|
time_diff = (image_times[-1] - image_times[0]).total_seconds()
|
|
period_minutes = time_diff / (len(image_times) - 1) / 60
|
|
result['sequence_period_minutes'] = round(period_minutes, 2)
|
|
|
|
# 添加时间段信息
|
|
result['time_range'] = {
|
|
'start': image_times[0].strftime("%Y-%m-%d %H:%M"),
|
|
'end': image_times[-1].strftime("%Y-%m-%d %H:%M")
|
|
}
|
|
|
|
save_result = self.db_handler.save_result(result)
|
|
print(f"Result saved to: {self.db_handler.results_collection.name}")
|
|
print(f"Result filename: {filename}")
|
|
return save_result
|
|
|
|
except Exception as e:
|
|
end_time = time.time()
|
|
processing_time = end_time - start_time
|
|
|
|
print(f"Error processing image sequence: {str(e)}")
|
|
print(f"Processing time (including error): {processing_time:.2f} seconds")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return False
|
|
# @staticmethod
|
|
# def extract_time_from_filename(filename):
|
|
# # 假设文件名格式为 "YYYYMMDDHHMMSS.jpg"
|
|
# time_str = filename.split('.')[0]
|
|
# return datetime.strptime(time_str, "%Y%m%d%H%M")
|
|
|
|
@staticmethod
|
|
def get_file_time(file_path):
|
|
# 获取文件的修改时间
|
|
mod_time = os.path.getmtime(file_path)
|
|
return datetime.fromtimestamp(mod_time)
|
|
|
|
def process_all_unprocessed_images(self, image_folders):
|
|
print(f"Searching for unprocessed images in: {image_folders}")
|
|
all_images = self.get_all_images(image_folders)
|
|
print(f"Found {len(all_images)} images in total")
|
|
selected_images = all_images[::10]
|
|
# Group images into sequences of 3
|
|
image_sequences = [selected_images[i:i+10] for i in range(0, len(selected_images), 10)]
|
|
# image_sequences = [all_images[i:i+10] for i in range(0, len(all_images), 10)]
|
|
|
|
processed_sequences = 0
|
|
for sequence in image_sequences:
|
|
if len(sequence) == 10:
|
|
self.process_image_sequence(sequence)
|
|
processed_sequences += 1
|
|
else:
|
|
print(f"Warning: Incomplete sequence. Found {len(sequence)} images.")
|
|
|
|
if processed_sequences == 0:
|
|
print("All current photos have been processed. Waiting for new photos...")
|
|
else:
|
|
print(f"Processed {processed_sequences} sequences.")
|
|
|
|
def run(self, root_folders):
|
|
print(f"Starting the system with root folder: {', '.join(root_folders)}")
|
|
|
|
while True:
|
|
current_time = datetime.now()
|
|
time_since_last_process = (current_time - self.last_processed_time).total_seconds()
|
|
|
|
if time_since_last_process >= 3600: # 1小时 = 3600秒
|
|
self.process_all_unprocessed_images(root_folders)
|
|
self.last_processed_time = current_time
|
|
|
|
# 计算下次检查的等待时间
|
|
wait_time = max(0, 3600 - (datetime.now() - self.last_processed_time).total_seconds())
|
|
print(f"Waiting for new photos... Next check in {wait_time:.0f} seconds.")
|
|
|
|
time.sleep(60) # 每分钟检查一次是否需要处理
|
|
# 使用示例
|
|
if __name__ == "__main__":
|
|
mongo_uri = "mongodb://minio_mongo:[email protected]:27017/minio_mongo"
|
|
db_name = "minio_mongo"
|
|
results_collection_name = "cpm"
|
|
|
|
model_dir = "worker_sys/OpenBMB/MiniCPM-V-2_6"
|
|
|
|
root_folders = [
|
|
"/www/wwwroot/zj.obscura.ac.cn/ipcam/Office/Cam2/CapturePics" ,
|
|
"/www/wwwroot/zj.obscura.ac.cn/ipcam/Office/Cam1/CapturePics"
|
|
] # 修改为 cam1 文件夹的路径
|
|
|
|
system = ImageSequenceAnalysisSystem(mongo_uri, db_name, model_dir, results_collection_name)
|
|
system.run(root_folders) |