m5stack unitv2 object detection and rectangles drawing

m5stack_unitv2_obj_detection.py

import subprocess
import json
import os
from datetime import datetime

import logging

import base64
from io import BytesIO
from PIL import Image, ImageDraw


LOGLEVEL = os.environ.get('LOGLEVEL', 'INFO').upper()
logging.basicConfig(level=LOGLEVEL, format="%(asctime)s %(message)s")

recognizer = subprocess.Popen(['/home/m5stack/payload/bin/object_recognition', '/home/m5stack/payload/uploads/models/yolo_20class'],
                                stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
recognizer.stdin.write("_{\"stream\":1}\r\n".encode('utf-8'))
recognizer.stdin.flush()

save_path = "/media/sdcard/imgs"
if not os.path.exists(save_path):
    os.makedirs(save_path)

save_next_image = False
obj_type = None
scale_x = 320 / 640
scale_y = 240 / 480

while True:
    line = recognizer.stdout.readline().decode('utf-8').strip()
    if not line:
        break
        
    #logging.info(line)
    try:
        doc = json.loads(line)
        
        if "obj" in doc:
            #logging.info(doc)
            current_objects = {(obj['x'] + round(obj['w'] / 2), obj['y'] + round(obj['h'] / 2)): obj for obj in doc.get('obj', []) if obj['type'] and obj['prob'] > 0.5}
            logging.info(current_objects)

            save_next_image=True
            
        if "img" in doc and save_next_image :

            image = Image.open(BytesIO(base64.b64decode(doc["img"])))

            draw = ImageDraw.Draw(image)

            for obj in current_objects.values():
                x, y, w, h = obj['x'], obj['y'], obj['w'], obj['h']

                x = int(obj['x'] * scale_x)
                y = int(obj['y'] * scale_y)
                w = int(obj['w'] * scale_x)
                h = int(obj['h'] * scale_y)
                obj_type=obj['type']
                draw.rectangle([x, y, x + w, y + h], outline="red", width=2)

            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            image_filename = os.path.join(save_path, f"detected_{obj_type}_{timestamp}.png")
            image.save(image_filename)
            logging.info(f"img saved {image_filename}")
            save_next_image=False
            
    except json.JSONDecodeError as e:
        logging.info(f"Error: Invalid JSON string\nJSONDecodeError: {str(e)}")