194 lines
7.3 KiB
Python
194 lines
7.3 KiB
Python
import torch
|
|
from PIL import Image
|
|
import json
|
|
from pymongo import MongoClient
|
|
import os
|
|
import glob
|
|
from datetime import datetime, timedelta
|
|
from ultralytics import YOLO
|
|
from bson import ObjectId
|
|
import time
|
|
|
|
# 数据库连接模块
|
|
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):
|
|
filename = result.get('filename', f"unknown_{result['timestamp']}")
|
|
|
|
existing_result = self.results_collection.find_one({'filename': filename})
|
|
if existing_result:
|
|
print(f"Image with filename {filename} has already been processed. Skipping.")
|
|
return False
|
|
|
|
result['filename'] = filename
|
|
|
|
if 'image_id' in result and isinstance(result['image_id'], ObjectId):
|
|
result['image_id'] = str(result['image_id'])
|
|
|
|
self.results_collection.insert_one(result)
|
|
return True
|
|
|
|
def is_image_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)
|
|
|
|
# YOLOv8nPoseProcessor 类
|
|
class YOLOv8nPoseProcessor:
|
|
def __init__(self, model_path):
|
|
self.model = YOLO(model_path)
|
|
|
|
def process_image(self, img):
|
|
results = self.model(img)
|
|
return results
|
|
|
|
def format_results(self, results):
|
|
result_data = []
|
|
for result in results:
|
|
boxes = result.boxes.xywh.tolist()
|
|
keypoints = result.keypoints.xy.tolist() if hasattr(result, 'keypoints') else None
|
|
classes = result.boxes.cls.tolist()
|
|
confs = result.boxes.conf.tolist()
|
|
|
|
for i, (box, cls, conf) in enumerate(zip(boxes, classes, confs)):
|
|
result_data.append({
|
|
'box': box,
|
|
'keypoints': keypoints[i] if keypoints else None,
|
|
'class': int(cls),
|
|
'class_name': self.model.names[int(cls)],
|
|
'confidence': float(conf)
|
|
})
|
|
|
|
if not result_data:
|
|
result_data.append({
|
|
'box': None,
|
|
'keypoints': None,
|
|
'class': None,
|
|
'class_name': None,
|
|
'confidence': None,
|
|
'message': 'No object detected'
|
|
})
|
|
|
|
return json.dumps(result_data)
|
|
|
|
# 主处理类
|
|
class ImageAnalysisSystem:
|
|
def __init__(self, mongo_uri, db_name, model_path, results_collection_name):
|
|
self.db_handler = DatabaseHandler(mongo_uri, db_name, results_collection_name)
|
|
self.image_processor = YOLOv8nPoseProcessor(model_path)
|
|
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
|
|
@staticmethod
|
|
def get_file_time(file_path):
|
|
# 获取文件的修改时间
|
|
mod_time = os.path.getmtime(file_path)
|
|
return datetime.fromtimestamp(mod_time)
|
|
def process_image(self, image_path):
|
|
print(f"Attempting to process image: {os.path.basename(image_path)}")
|
|
try:
|
|
# json_folder = os.path.join("/www/wwwroot/zj.obscura.ac.cn/ipcam/Office/Cam2", 'json')
|
|
# json_filename = f"{os.path.basename(image_path).split('.')[0]}.json"
|
|
# json_path = os.path.join(json_folder, json_filename)
|
|
|
|
# if os.path.exists(json_path):
|
|
# print(f"Skipping already processed image: {json_path}")
|
|
# return
|
|
|
|
filename = os.path.basename(image_path)
|
|
|
|
if self.db_handler.is_image_processed(filename):
|
|
print(f"Skipping already processed image: {filename}")
|
|
return False
|
|
|
|
print("Processing new image...")
|
|
|
|
image = Image.open(image_path)
|
|
results = self.image_processor.process_image(image)
|
|
formatted_results = self.image_processor.format_results(results)
|
|
|
|
# timestamp = datetime.now()
|
|
file_timestamp = self.get_file_time(image_path)
|
|
result = {
|
|
'timestamp': file_timestamp.strftime("%Y%m%d_%H%M%S"),
|
|
'image_path': image_path,
|
|
'filename': os.path.basename(image_path),
|
|
'results': json.loads(formatted_results)
|
|
}
|
|
|
|
# os.makedirs(json_folder, exist_ok=True)
|
|
# with open(json_path, 'w', encoding='utf-8') as f:
|
|
# json.dump(result, f, ensure_ascii=False, indent=4, cls=JSONEncoder)
|
|
|
|
# self.db_handler.save_result(result)
|
|
# print(f"Processed image at: {timestamp}")
|
|
# print(f"JSON saved to: {json_path}")
|
|
# print(f"result saved to: {results_collection_name}")
|
|
if self.db_handler.save_result(result):
|
|
print(f"Result saved to: {self.db_handler.results_collection.name}")
|
|
return True
|
|
else:
|
|
print(f"Image {filename} was already in the database. Skipping.")
|
|
return False
|
|
|
|
except Exception as e:
|
|
print(f"Error processing image: {str(e)}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
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")
|
|
|
|
processed_count = 0
|
|
for image_path in all_images:
|
|
if self.process_image(image_path):
|
|
processed_count += 1
|
|
|
|
return processed_count
|
|
|
|
def run(self, root_folder):
|
|
print(f"Starting the system with root folder: {', '.join(root_folder)}")
|
|
# image_folder = os.path.join(root_folder)
|
|
|
|
while True:
|
|
print("Checking for unprocessed images...")
|
|
processed_count = self.process_all_unprocessed_images(root_folder)
|
|
|
|
if processed_count > 0:
|
|
print(f"Finished processing {processed_count} images.")
|
|
else:
|
|
print("No new images to process. Waiting for new images...")
|
|
|
|
# 等待一段时间后再次检查新图片
|
|
time.sleep(60) # 每分钟检查一次是否有新图片
|
|
|
|
# 使用示例
|
|
if __name__ == "__main__":
|
|
mongo_uri = "mongodb://minio_mongo:[email protected]:27017/minio_mongo"
|
|
db_name = "minio_mongo"
|
|
results_collection_name = "pose"
|
|
|
|
model_path = "worker_sys/function/yolov8x-pose.pt" # 请确保这个路径指向你的YOLO-Pose模型文件
|
|
|
|
root_folder = [
|
|
"/www/wwwroot/zj.obscura.ac.cn/ipcam/Office/Cam2/CapturePics" ,
|
|
"/www/wwwroot/zj.obscura.ac.cn/ipcam/Office/Cam1/CapturePics"
|
|
] # 修改为 cam1 文件夹的路径
|
|
|
|
system = ImageAnalysisSystem(mongo_uri, db_name, model_path, results_collection_name)
|
|
system.run(root_folder) |