zydi-api/api_history/function/producer-minio.py

import json
import yaml
from kafka import KafkaConsumer, KafkaProducer
import time
from datetime import datetime
import redis
from minio import Minio
import io
from PIL import Image
import base64
import traceback


# 加载配置文件
with open('worker_sys/function/config.yaml', 'r') as file:
    config = yaml.safe_load(file)

# Kafka 配置
kafka_config = {
    'bootstrap_servers': config['kafka']['bootstrap_servers'],
    'value_serializer': lambda v: json.dumps(v).encode('utf-8') if config['kafka']['value_serializer'] == 'json' else None
}

# 创建Kafka生产者
producer = KafkaProducer(**kafka_config)

# 创建Kafka消费者（用于输入主题）
consumer = KafkaConsumer(
    config['kafka']['input_topic']['name'],
    group_id='image-processor',
    bootstrap_servers=config['kafka']['bootstrap_servers']
)

# Redis 配置
redis_config = config['redis']

# 创建 Redis 客户端
redis_client = redis.Redis(**redis_config)

# MinIO 配置
minio_config = config['minio']

# 创建 MinIO 客户端
minio_client = Minio(
    minio_config['endpoint'],
    access_key=minio_config['access_key'],
    secret_key=minio_config['secret_key'],
    secure=minio_config['secure']
)

# Kafka topics
topic_all_frames = config['kafka']['topics']['all_frames']['name']
topic_ten_seconds = config['kafka']['topics']['ten_seconds']['name']
topic_input = config['kafka']['input_topic']['name']

# 消费者数量
NUM_CONSUMERS_ALL_FRAMES = config['kafka']['topics']['all_frames']['num_consumers']
NUM_CONSUMERS_TEN_SECONDS = config['kafka']['topics']['ten_seconds']['num_consumers']
NUM_CONSUMERS_INPUT = config['kafka']['input_topic']['num_consumers']

def parse_key_name(key_name):
    parts = key_name.split('/')
    bucket = parts[0]
    image_name = '/'.join(parts[1:])
    image_parts = image_name.rsplit('.', 1)[0].split('_')
    camera_name = image_parts[0]
    timestamp = '_'.join(image_parts[1:3])
    return bucket, camera_name, timestamp

def should_send_to_ten_seconds_topic(timestamp):
    try:
        dt = datetime.strptime(timestamp, '%Y%m%d_%H%M%S')
        return dt.second % 10 == 0
    except ValueError as e:
        print(f"Error parsing timestamp '{timestamp}': {str(e)}")
        return False

def get_image_from_minio_and_cache(key_name):
    bucket, object_name = key_name.split('/', 1)
    cache_key = f"{key_name}"
    
    try:
        response = minio_client.get_object(bucket, object_name)
        image_data = response.read()
        print(f"Successfully retrieved image from MinIO: {bucket}/{object_name}")

        image = Image.open(io.BytesIO(image_data))
        buffered = io.BytesIO()
        image.save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue()).decode()

        redis_client.setex(cache_key, 86400, img_str)
        print(f"Successfully cached image in Redis: {cache_key}")
        
        return cache_key
    except Exception as e:
        print(f"Error in get_image_from_minio_and_cache: {str(e)}")
        print("Traceback:")
        traceback.print_exc()
        return None

def process_message(message):

    value_data = json.loads(message.value.decode('utf-8'))
    key_name = value_data['Key']
    print(f"Processing key: {key_name}")
    
    # 从 key_name 中提取 bucket
    parts = key_name.split('/', 1)
    if len(parts) != 2:
        print(f"Error: Invalid key format. Expected 'bucket/object', got '{key_name}'")
        return
    
    bucket, object_name = parts
    
    # 解析 camera_name 和 timestamp
    camera_name, timestamp = parse_object_name(object_name)
    
    cache_key = get_image_from_minio_and_cache(key_name)
    
    if cache_key is None:
        print(f"Failed to process image: {key_name}")
        return
    # 提取 etag 和 size 信息
    object_info = value_data['Records'][0]['s3']['object']
    etag = object_info.get('eTag', '')
    size = object_info.get('size', 0)


    message_data = {
        'bucket': bucket,
        'camera_name': camera_name,
        'timestamp': timestamp,
        'object': object_name,
        'cache_key': cache_key,
        'etag': etag,
        'size': size
    }
    
    # 发送到 pose-input 主题
    producer.send(topic_all_frames, value=message_data)
    print(f"Sent message to {topic_all_frames}: {key_name}")

    # 发送到 cpm-input 主题
    producer.send(topic_ten_seconds, value=message_data)
    print(f"Sent message to {topic_ten_seconds}: {key_name}")

    # #只有在满足特定条件时才发送到 cpm-input 主题
    # if should_send_to_ten_seconds_topic(timestamp):
    #     producer.send(TOPIC_TEN_SECONDS, value=message_data)
    #     print(f"Sent message to {TOPIC_TEN_SECONDS}: {key_name}")

    producer.flush()
    print(f"Successfully processed and sent messages for: {key_name}")


def parse_object_name(object_name):
    # 假设对象名格式为 "cameraX_YYYYMMDD_HHMMSS.jpg"
    parts = object_name.split('_')
    if len(parts) != 3:
        raise ValueError(f"Invalid object name format: {object_name}")
    
    camera_name = parts[0]
    timestamp = f"{parts[1]}_{parts[2].split('.')[0]}"
    return camera_name, timestamp

def get_image_from_minio_and_cache(key_name):
    print(f"Received key_name: {key_name}")
    
    try:
        bucket, object_name = key_name.split('/', 1)
    except ValueError as e:
        print(f"Error splitting key_name '{key_name}': {str(e)}")
        return None
    
    cache_key = f"{key_name}"
    
    try:
        response = minio_client.get_object(bucket, object_name)
        image_data = response.read()
        print(f"Successfully retrieved image from MinIO: {bucket}/{object_name}")

        image = Image.open(io.BytesIO(image_data))
        buffered = io.BytesIO()
        image.save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue()).decode()

        redis_client.setex(cache_key, 86400, img_str)
        print(f"Successfully cached image in Redis: {cache_key}")
        
        return cache_key
    except Exception as e:
        print(f"Error in get_image_from_minio_and_cache: {str(e)}")
        print("Traceback:")
        traceback.print_exc()
        return None

def main():
    print("Starting image processing service...")
    for message in consumer:
        process_message(message)

if __name__ == "__main__":
    main()