Zordon/lib/render_workers/blender_worker.py

#!/usr/bin/env python3
import json
import re
from .render_worker import *


class Blender(BaseRenderEngine):

    install_paths = ['/Applications/Blender.app/Contents/MacOS/Blender']
    supported_extensions = ['.blend']

    @classmethod
    def version(cls):
        version = None
        try:
            render_path = cls.renderer_path()
            if render_path:
                ver_out = subprocess.check_output([render_path, '-v'])
                version = ver_out.decode('utf-8').splitlines()[0].replace('Blender', '').strip()
        except Exception as e:
            logging.error(f'Failed to get Blender version: {e}')
        return version

    @classmethod
    def get_output_formats(cls):
        format_string = cls.get_help().split('Format Options')[-1].split('Animation Playback Options')[0]
        formats = re.findall(r"'([A-Z_0-9]+)'", format_string)
        return formats

    @classmethod
    def full_report(cls):
        return {'version': cls.version(),
                'help_text': cls.get_help(),
                'formats': cls.get_formats()}

    @classmethod
    def run_python_expression(cls, project_path, python_expression):
        if os.path.exists(project_path):
            try:
                return subprocess.run([cls.renderer_path(), '-b', project_path, '--python-expr', python_expression],
                                      capture_output=True)
            except Exception as e:
                logger.warning(f"Error running python expression in blender: {e}")
                pass
        else:
            raise FileNotFoundError(f'Project file not found: {project_path}')

    @classmethod
    def run_python_script(cls, project_path, script_path):
        if os.path.exists(project_path) and os.path.exists(script_path):
            try:
                return subprocess.run([cls.renderer_path(), '-b', project_path, '--python', script_path],
                                      capture_output=True)
            except Exception as e:
                logger.warning(f"Error running python expression in blender: {e}")
                pass
        elif not os.path.exists(project_path):
            raise FileNotFoundError(f'Project file not found: {project_path}')
        elif not os.path.exists(script_path):
            raise FileNotFoundError(f'Python script not found: {script_path}')
        raise Exception("Uncaught exception")

    @classmethod
    def get_scene_info(cls, project_path):
        scene_info = None
        try:
            results = cls.run_python_script(project_path, os.path.join(os.path.dirname(os.path.realpath(__file__)),
                                                                   'scripts', 'get_blender_info.py'))
            result_text = results.stdout.decode()
            for line in result_text.splitlines():
                if line.startswith('SCENE_DATA:'):
                    raw_data = line.split('SCENE_DATA:')[-1]
                    scene_info = json.loads(raw_data)
                    break
        except Exception as e:
            logger.error(f'Error getting file details for .blend file: {e}')
        return scene_info

    @classmethod
    def pack_project_file(cls, project_path):
        # Credit to L0Lock for pack script - https://blender.stackexchange.com/a/243935
        pack_expression = "import bpy\n" \
                      "bpy.ops.file.pack_all()\n" \
                      "myPath = bpy.data.filepath\n" \
                      "myPath = str(myPath)\n" \
                      "bpy.ops.wm.save_as_mainfile(filepath=myPath[:-6]+'_packed'+myPath[-6:])"

        try:
            results = Blender.run_python_expression(project_path, pack_expression)

            result_text = results.stdout.decode()
            dir_name = os.path.dirname(project_path)

            # report any missing textures
            not_found = re.findall("(Unable to pack file, source path .*)\n", result_text)
            for err in not_found:
                logger.error(err)

            p = re.compile('Info: Saved "(.*)"')
            match = p.search(result_text)
            if match:
                new_path = os.path.join(dir_name, match.group(1))
                logger.info(f'Blender file packed successfully to {new_path}')
                return new_path
        except Exception as e:
            logger.error(f'Error packing .blend file: {e}')
        return None


class BlenderRenderWorker(BaseRenderWorker):

    engine = Blender

    def __init__(self, input_path, output_path, args=None):
        super(BlenderRenderWorker, self).__init__(input_path=input_path, output_path=output_path,
                                                  ignore_extensions=False, args=args)

        # 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)
        self.render_all_frames = self.args.get('render_all_frames', False) or \
                                 '-a' in (self.args.get('raw', None) or "").split(' ')
        self.frame_to_render = 0

        # Stats
        self.__frame_percent_complete = 0.0

        # Scene Info
        self.scene_info = Blender.get_scene_info(input_path)
        self.total_frames = (int(self.scene_info.get('frame_end', 0)) - int(self.scene_info.get('frame_start', 0)) + 1) \
            if self.render_all_frames else 1
        self.current_frame = int(self.scene_info.get('frame_start', 0))

    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}'];"])

        cmd.extend(['-E', self.blender_engine, '-o', self.output_path, '-F', self.export_format])

        # all frames or single
        cmd.extend(['-a'] if self.render_all_frames else ['-f', str(self.frame_to_render)])

        # 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 'error' in line.lower():
            logger.error(line)
            self.errors.append(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]
                if self.render_all_frames:
                    logger.debug(f'Frame {self.current_frame} completed in {time_completed}')
                else:
                    logger.info(f'Render completed in {time_completed}')
            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)


if __name__ == '__main__':

    print(Blender.full_report())

    x = Blender.pack_project_file('/Users/brettwilliams/Downloads/barbershop_interior.blend')
    # print(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)