#!/usr/bin/env python3

import re
import rospy
import lasr_vision_bodypix as bodypix

from sensor_msgs.msg import Image
from lasr_vision_msgs.srv import BodyPixDetection, BodyPixDetectionRequest, BodyPixDetectionResponse

# Initialise rospy
rospy.init_node('bodypix_service')

# Determine variables
DEBUG = rospy.get_param('~debug', False)
PRELOAD = rospy.get_param('~preload', []) # resnet50 or mobilenet50

for model in PRELOAD:
    pass

# Keep track of publishers
debug_publishers = {}

def detect(request: BodyPixDetectionRequest) -> BodyPixDetectionResponse:
    '''
    Hand off detection request to bodypix library
    '''
    debug_publisher = None
    if DEBUG:
        if request.dataset in debug_publishers:
            debug_publisher = debug_publishers[request.dataset]
        else:
            topic_name = re.sub(r'[\W_]+', '', request.dataset)
            debug_publisher = rospy.Publisher(f'/bodypix/debug/{topic_name}', Image, queue_size=1)
    return bodypix.detect(request, debug_publisher)

rospy.Service('/bodypix/detect', BodyPixDetection, detect)
rospy.loginfo('BodyPix service started')
rospy.spin()
