Zordon/lib/workers/blender_worker.py

#!/usr/bin/env python3
import json
import re
try:
    from .base_worker import *
except ImportError:
    from base_worker import *

from ..engines.blender_engine import Blender


class BlenderRenderWorker(BaseRenderWorker):

    engine = Blender

    def __init__(self, input_path, output_path, priority=2, args=None, owner=None,
                      client=None, name=None):
        super(BlenderRenderWorker, self).__init__(input_path=input_path, output_path=output_path, args=args,
                                                 client=client, priority=priority, owner=owner, name=name)

        # Args
        self.blender_engine = self.args.get('engine', 'BLENDER_EEVEE').upper()
        self.export_format = self.args.get('export_format', None) or 'JPEG'
        self.camera = self.args.get('camera', None)

        # Stats
        self.__frame_percent_complete = 0.0

        # Scene Info
        self.scene_info = Blender.get_scene_info(input_path)
        self.start_frame = int(self.scene_info.get('start_frame', 1))
        self.end_frame = int(self.scene_info.get('end_frame', self.start_frame))
        self.project_length = (self.end_frame - self.start_frame) + 1
        self.current_frame = -1

    def generate_worker_subprocess(self):

        cmd = [self.engine.renderer_path()]
        if self.args.get('background', True):  # optionally run render not in background
            cmd.append('-b')
        cmd.append(self.input_path)

        if self.camera:
            cmd.extend(['--python-expr', f"import bpy;bpy.context.scene.camera = bpy.data.objects['{self.camera}'];"])

        # add dash at end of given path to separate frame numbers
        path_with_ending_dash = os.path.splitext(self.output_path)[0] + "-" + os.path.splitext(self.output_path)[1]
        cmd.extend(['-E', self.blender_engine, '-o', path_with_ending_dash, '-F', self.export_format])

        # set frame range
        cmd.extend(['-s', self.start_frame, '-e', self.end_frame, '-a'])

        # Convert raw args from string if available
        raw_args = self.get_raw_args()
        if raw_args:
            cmd.extend(raw_args)

        return cmd

    def _parse_stdout(self, line):

        pattern = re.compile(
            r'Fra:(?P<frame>\d*).*Mem:(?P<memory>\S+).*Time:(?P<time>\S+)(?:.*Remaining:)?(?P<remaining>\S*)')
        found = pattern.search(line)
        if found:
            stats = found.groupdict()
            memory_use = stats['memory']
            time_elapsed = stats['time']
            time_remaining = stats['remaining'] or 'Unknown'

            sample_string = line.split('|')[-1].strip()
            if "sample" in sample_string.lower():
                samples = re.sub(r'[^\d/]', '', sample_string)
                self.__frame_percent_complete = int(samples.split('/')[0]) / int(samples.split('/')[-1])

            # Calculate rough percent based on cycles
            # EEVEE
            # 10-Apr-22 22:42:06 - RENDERER: Fra:0 Mem:857.99M (Peak 928.55M) | Time:00:03.96 | Rendering 1 / 65 samples
            # 10-Apr-22 22:42:10 - RENDERER: Fra:0 Mem:827.09M (Peak 928.55M) | Time:00:07.92 | Rendering 26 / 64 samples
            # 10-Apr-22 22:42:10 - RENDERER: Fra:0 Mem:827.09M (Peak 928.55M) | Time:00:08.17 | Rendering 51 / 64 samples
            # 10-Apr-22 22:42:10 - RENDERER: Fra:0 Mem:827.09M (Peak 928.55M) | Time:00:08.31 | Rendering 64 / 64 samples

            # CYCLES
            # 10-Apr-22 22:43:22 - RENDERER: Fra:0 Mem:836.30M (Peak 1726.13M) | Time:00:01.56 | Remaining:00:30.65 | Mem:588.68M, Peak:588.68M | Scene, View Layer | Sample 1/150
            # 10-Apr-22 22:43:43 - RENDERER: Fra:0 Mem:836.30M (Peak 1726.13M) | Time:00:22.01 | Remaining:00:03.36 | Mem:588.68M, Peak:588.68M | Scene, View Layer | Sample 129/150

            if int(stats['frame']) > self.current_frame:
                self.current_frame = int(stats['frame'])
                logger.debug(
                    'Frame:{0} | Mem:{1} | Time:{2} | Remaining:{3}'.format(self.current_frame, memory_use,
                                                                            time_elapsed, time_remaining))
        elif "file doesn't exist" in line.lower():
            self.log_error(line, halt_render=True)
        elif line.lower().startswith('error'):
            self.log_error(line)
        elif 'Saved' in line or 'Saving' in line or 'quit' in line:
            match = re.match(r'Time: (.*) \(Saving', line)
            if match:
                time_completed = match.groups()[0]
                frame_count = self.current_frame - self.end_frame + self.total_frames
                logger.info(f'Frame #{self.current_frame} - '
                            f'{frame_count} of {self.total_frames} completed in {time_completed} | '
                            f'Total Elapsed Time: {datetime.now() - self.start_time}')
            else:
                logger.debug(line)
        else:
            pass
            # if len(line.strip()):
            #     logger.debug(line.strip())

    def percent_complete(self):
        if self.total_frames <= 1:
            return self.__frame_percent_complete
        else:
            whole_frame_percent = (self.current_frame - 1) / self.total_frames
            adjusted_frame_percent = self.__frame_percent_complete / self.total_frames
            total_percent = whole_frame_percent + adjusted_frame_percent
            return max(total_percent, 0)

    def post_processing(self):
        output_dir = os.listdir(os.path.dirname(self.output_path))
        if self.total_frames > 1 and len(output_dir) > 1:
            from ..utilities.ffmpeg_helper import image_sequence_to_video
            logger.info("Generating preview for image sequence")

            # get proper file extension
            path_with_ending_dash = os.path.splitext(self.output_path)[0] + "-"
            found_output = next(obj for obj in output_dir if os.path.basename(path_with_ending_dash) in obj)
            glob_pattern = path_with_ending_dash + '%04d' + ('.' + found_output.split('.')[-1] if found_output else "")

            try:
                image_sequence_to_video(source_glob_pattern=glob_pattern,
                                        output_path=self.output_path + '.mp4',
                                        framerate=self.scene_info['fps'])
                logger.info('Successfully generated preview video from image sequence')
            except Exception as e:
                logger.error(f'Error generating video from image sequence: {e}')


if __name__ == '__main__':

    import pprint
    x = Blender.get_scene_info('/Users/brett/Desktop/TC Previz/nallie_farm.blend')
    pprint.pprint(x)

    # logging.basicConfig(format='%(asctime)s - %(message)s', datefmt='%d-%b-%y %H:%M:%S', level=logging.DEBUG)
    # r = BlenderRenderWorker('/Users/brett/Blender Files/temple_animatic.blend', '/Users/brett/testing1234')
    # # r.engine = 'CYCLES'
    # r.start()
    # while r.is_running():
    #     time.sleep(1)