blender_poly.py

bl_info = {
    "name": "Blender Poly",
    "category": "Object",
    "author": "Yuichi Sato",
    "version": (2, 0),
    "blender": (2, 80, 0),
    "location": "Object Panel > Poly",
    "wiki_url": "https://github.com/satoyuichi/BlenderPoly",
}

import bpy
import bpy.utils.previews
import requests
import json
import re
import os
import tempfile
import sys
from pathlib import Path

__package__ = "blender_poly"
BLENDER_POLY_PATH = 'BlenderPoly'

preview_collections = {}
blender_poly_json = {}
blender_poly_category_items = [
    ('animals', 'Animals and Creatures', 'animals'),
    ('architecture', 'Architecture', 'architecture'),
    ('art', 'Art', 'art'),
    ('current_events', 'Current events', 'current_events'),
    ('food', 'Food and Drink', 'food'),
    ('furniture_home', 'Furniture and Home', 'furniture_home'),
    ('history', 'history', 'history'),
    ('nature', 'Nature', 'nature'),
    ('objects', 'Objects', 'objects'),
    ('people', 'People and Characters', 'people'),
    ('scenes', 'Places and Scenes', 'scenes'),
    ('science', 'Science', 'science'),
    ('sports_fitness', 'Sports and Fitness', 'sports_fitness'),
    ('tech', 'Technology', 'tech'),
    ('transport', 'Transport', 'transport'),
    ('travel', 'Travel', 'travel')
]

def get_temp_path (context):
    props = context.window_manager.poly
    return Path (tempfile.gettempdir()).joinpath (BLENDER_POLY_PATH, props.category_type)

def get_element_from_json (id):
    global blender_poly_json

    if not 'assets' in blender_poly_json.keys ():
        return None

    for elem in blender_poly_json['assets']:
        if elem['name'] == id:
            return elem

def enum_previews_from_model_previews(self, context):
    """EnumProperty callback"""
    if context is None:
        return []

    wm = context.window_manager
    preferences = context.preferences.addons[__package__].preferences
    props = context.window_manager.poly
    
    enum_items = []
    directory = get_temp_path (context)

    pcoll = preview_collections[props.category_type]
    
    if directory == pcoll.previews_previews_dir:
        return pcoll.previews_previews
        
    filepath_list = list (Path (directory).glob ('**/*'))
    
    # Load JSON file.
    json_path = directory.joinpath (props.category_type + ".json")
    if not json_path.exists ():
        return enum_items
    
    with json_path.open ("r", encoding='utf-8') as f:
        global blender_poly_json
        blender_poly_json = json.loads (f.read ())
    
    for i, filepath in enumerate (filepath_list):
        if filepath.suffix != ".png":
            continue

        # Load image.
        comp_path = str (filepath)
        thumb = pcoll.load (comp_path, comp_path, 'IMAGE')

        # Get asset name from JSON.
        id_name = 'assets/' + filepath.stem
        elem = get_element_from_json (id_name)

        enum_items.append ((id_name, elem['displayName'], comp_path, thumb.icon_id, i))
    
    pcoll.previews_previews = enum_items
    pcoll.previews_previews_dir = directory
    return pcoll.previews_previews

def import_mtl_by_url(context, url, relativePath):
    r = requests.get(url)
    
    print (url)

    file_path = get_temp_path (context).joinpath (relativePath)
    with file_path.open ("w", encoding='utf-8') as f:
        f.write (r.text)

def import_obj_by_url(context, url, relativePath):
    r = requests.get(url)
    
    print (url)

    file_path = get_temp_path (context).joinpath (relativePath)
    with file_path.open ("w", encoding='utf-8') as f:
        f.write (r.text)

    # Create empty and initialize
    bpy.ops.object.empty_add()
    empty = context.active_object
    empty.name = get_element_from_json (context.window_manager.poly_model_previews)['displayName']
    empty.show_name = True
    
    # Import obj
    bpy.ops.import_scene.obj(filepath=str(file_path), axis_forward='-Z', axis_up='Y', filter_glob="*.obj;*.mtl", use_edges=True, use_smooth_groups=True, use_split_objects=True, use_split_groups=True, use_groups_as_vgroups=False, use_image_search=True, split_mode='ON', global_clight_size=0)
    bpy.ops.object.transform_apply(rotation=True)

    # Set empty location
#        props = context.window_manager.poly.emptyPivot
    emptyPivot = context.window_manager.poly.emptyPivot

    # Initialize
    empty.location[0] = 0.0
    empty.location[1] = 0.0
    empty.location[2] = 0.0
    pivot_index = 2
    pivot_coord = 0.0

    if emptyPivot == 'x_min' or emptyPivot == 'x_max':
        pivot_index = 0
    elif emptyPivot == 'y_min' or emptyPivot == 'y_max':
        pivot_index = 1
    else:
        pivot_index = 2

    if emptyPivot == 'x_min' or emptyPivot == 'y_min' or emptyPivot == 'z_min':
        pivot_coord = sys.maxsize
        for child in context.selected_objects:
            for coord in child.bound_box:
                pivot_coord = coord[pivot_index] if (pivot_coord > coord[pivot_index]) else pivot_coord
    else:
        pivot_coord = -sys.maxsize
        for child in context.selected_objects:
            for coord in child.bound_box:
                pivot_coord = coord[pivot_index] if (pivot_coord < coord[pivot_index]) else pivot_coord
    
    empty.location[pivot_index] = pivot_coord

    bpy.ops.object.parent_set()
    
    # Set active only empty
    bpy.ops.object.select_all(action='DESELECT')
    empty.select_set(True)
    
def import_obj_and_mtl(context, json):
    # Load mtl
    url = json['resources'][0]['url']
    relativePath = json['resources'][0]['relativePath']
    import_mtl_by_url (context, url, relativePath)
    
    # Load obj
    url = json['root']['url']
    relativePath = json['root']['relativePath']
    import_obj_by_url (context, url, relativePath)
        
class BPLY_Preferences(bpy.types.AddonPreferences):
    bl_idname = __package__

    polyApiKey: bpy.props.StringProperty (
        name = "API Key",
        default = "",
        description = "Input the Poly's API Key",
        subtype = 'BYTE_STRING'
        )
           
    def draw(self, context):
        props = context.window_manager.poly
        layout = self.layout
        preferences = context.preferences.addons[__package__].preferences
        
        row = layout.row()
        row.prop(preferences, 'polyApiKey')
        
        row = layout.row()
        row.scale_y = 1.25
        row.operator("scene.poly_install_assets", icon='PREFERENCES')

class BPLY_OT_InstallAssets(bpy.types.Operator):
    """Save the Blender Poly assets filepath"""
    bl_idname = "scene.poly_install_assets"
    bl_label = "Save Settings"
    bl_options = {'REGISTER', 'UNDO'}

    def execute(self, context):
        bpy.ops.wm.save_userpref()

        return {'FINISHED'}
    
class BlenderPolyProps(bpy.types.PropertyGroup):
    # Select empty pivot
    emptyPivot: bpy.props.EnumProperty(
        items=[
            ('x_min', 'X Min', 'X Min'),
            ('x_max', 'X Max', 'X Max'),
            ('y_min', 'Y Min', 'Y Min'),
            ('y_max', 'Y Max', 'Y Max'),
            ('z_min', 'Z Min', 'Z Min'),
            ('z_max', 'Z Max', 'Z Max'),
        ],
        name="Empty pivot",
        default="z_min")
    
    # Normal import
    category_type: bpy.props.EnumProperty(
        items=blender_poly_category_items,
        name="Category type",
        default="animals")
    maxComplexity: bpy.props.EnumProperty(
        items=[
            ('COMPLEX', 'COMPLEX', 'COMPLEX'),
            ('MEDIUM', 'MEDIUM', 'MEDIUM'),
            ('SIMPLE', 'SIMPLE', 'SIMPLE')
        ],
        name="Max complexity",
        default="COMPLEX")
    keywords: bpy.props.StringProperty(name='Keywords', description='Keywords')
    curated: bpy.props.BoolProperty(name='Curated', description='Curated')
    pageSize: bpy.props.IntProperty(name='Page size', description='Page size', default=20, min=1, max=100)
    orderBy: bpy.props.EnumProperty(
        items=[('BEST', 'BEST', 'BEST'), ('NEWEST', 'NEWEST', 'NEWEST'), ('OLDEST', 'OLDEST', 'OLDEST')],
        name='Order by',
        default='BEST')
    nextPageToken: bpy.props.StringProperty(name='nextPageToken', default='', description='Token')
    
    # Direct import
    directID: bpy.props.StringProperty(name='ID', description='Import model ID')
        
class BPLY_PT_LayoutPanel(bpy.types.Panel):
    """Creates a Panel in the scene context of the properties editor"""
    bl_label = "Blender Poly"
    bl_space_type = 'VIEW_3D'
    bl_region_type = 'UI'
    bl_category = "Create"
    bl_context = "objectmode"

    def draw(self, context):
        props = context.window_manager.poly
        
        layout = self.layout
        wm = context.window_manager

        scene = context.scene

        box = layout.box()
        
        row = box.row(align=True)
        row.prop(props, "emptyPivot")
        
        # Normal import
        box = layout.box()
        
        row = box.row(align=True)
        row.prop(props, "category_type")
        
        row = box.row(align=True)
        row.prop(props, "maxComplexity")

        row = box.row(align=True)
        row.prop(props, "orderBy")

        row = box.row(align=True)
        row.prop(props, "pageSize")
        
        row = box.row(align=True)
        row.prop(props, "curated")
        
        row = box.row(align=True)
        row.prop(props, "keywords")
        
        row = box.row(align=True)
        col = row.column()
        col.scale_y = 2.0
        if props.nextPageToken:
            col.operator("blender_poly.to_head", text="Head", icon="REW")

            col = row.column()
            col.scale_y = 2.0
            col.operator("blender_poly.load", text="Next", icon="FORWARD")
        else:
            col.operator("blender_poly.load", text="Load", icon="FILE_FOLDER")

        row = box.row(align=True)
        col = row.column()
        col.scale_y = 1
        col.template_icon_view(wm, "poly_model_previews", show_labels=True)
        elem = get_element_from_json (context.window_manager.poly_model_previews)
        if elem == None:
            col.label(text='')
        else:
            col.label(text=elem['displayName'])
        
        row = box.row(align=True)
        row.scale_y = 1.5
        row.operator("blender_poly.import", text="Import", icon="IMPORT")

        # Direct import
        box = layout.box()
        
        row = box.row(align=True)
        row.label(text="Direct import:")
        
        row = box.row(align=True)
        row.prop(props, "directID")
        row = box.row(align=True)
        row.operator("blender_poly.direct_import", text="Direct Import", icon="IMPORT")

class BPLY_OT_ToHead(bpy.types.Operator):
    bl_idname = "blender_poly.to_head"
    bl_label = "To Head Operator"

    def execute(self, context):
        props = context.window_manager.poly
        props.nextPageToken = ''
        return {'FINISHED'}
    
class BPLY_OT_DirectImport(bpy.types.Operator):
    bl_idname = "blender_poly.direct_import"
    bl_label = "Direct Import Operator"

    def execute(self, context):
        props = context.window_manager.poly
        preferences = context.preferences.addons[__package__].preferences
        
        payload = { 'key': preferences.polyApiKey }

        r = requests.get ("https://poly.googleapis.com/v1/assets/" + props.directID, params=payload)
        json = r.json ()
        
        for elem in json['formats']:
            if (elem['formatType']) == 'OBJ':
                import_obj_and_mtl (context, elem)
                break

        return {'FINISHED'}
            
class BPLY_OT_AssetsLoader(bpy.types.Operator):
    bl_idname = "blender_poly.load"
    bl_label = "Load Operator"
    
    def getPayload (self, preferences, props):
        return {'key': preferences.polyApiKey, 'format': 'OBJ',
            'category': props.category_type,
            'maxComplexity': props.maxComplexity,
            'curated': props.curated,
            'keywords': props.keywords,
            'pageSize': props.pageSize,
            'orderBy': props.orderBy,
            'pageToken': props.nextPageToken
        }
        
    def writeThumbnails (self, json, thumbnail_path):
        for asset in json['assets']:
            suffix = Path (asset['thumbnail']['relativePath']).suffix
            asset['name'] = re.sub (r'assets/', '', asset['name'])
            filepath = thumbnail_path.joinpath (asset['name']).with_suffix(suffix)

            if not filepath.exists ():
                thumbnail = requests.get (asset['thumbnail']['url'])
                with filepath.open (mode='wb') as f:
                    f.write (thumbnail.content)

    def recreatePreviews (self, props):
        pcoll = preview_collections[props.category_type]
        bpy.utils.previews.remove(pcoll)
        pcoll = bpy.utils.previews.new ()
        preview_collections[props.category_type] = pcoll
        pcoll.previews_previews = ()
        pcoll.previews_previews_dir = ""

    def execute(self, context):
        if not __package__ in context.preferences.addons.keys ():
            return {'FINISHED'}       

        tmp_path = get_temp_path (context)
        if tmp_path.exists ():
            filepath_list = list (tmp_path.glob ('**/*'))
            for path in filepath_list:
                os.remove (str(path))
        else:
            tmp_path.mkdir (parents=True)

        props = context.window_manager.poly
        preferences = context.preferences.addons[__package__].preferences
        payload = self.getPayload(preferences, props)

        self.recreatePreviews (props)

        r = requests.get ("https://poly.googleapis.com/v1/assets", params=payload)
        json = r.json()

        if not 'assets' in json.keys ():
            return {'INTERFACE'}

        props.nextPageToken = json['nextPageToken']

        # Save JSON
        json_path = tmp_path.joinpath (props.category_type + ".json")
        with json_path.open ("w", encoding='utf-8') as f:
            f.write (r.text)

        self.writeThumbnails (json, tmp_path)

        return {'FINISHED'}

class BPLY_OT_AssetsImporter(bpy.types.Operator):
    bl_idname = "blender_poly.import"
    bl_label = "Import Operator"

    def execute(self, context):
        global blender_poly_json

        elem = get_element_from_json (context.window_manager.poly_model_previews)
        
        obj_elem = ''
        for el in elem['formats']:
            if el['formatType'] == 'OBJ':
                obj_elem = el

        import_obj_and_mtl(context, obj_elem)
        
        return {'FINISHED'}
        
def register():
    bpy.utils.register_class(BlenderPolyProps)
    bpy.utils.register_class(BPLY_PT_LayoutPanel)
    bpy.utils.register_class(BPLY_OT_DirectImport)
    bpy.utils.register_class(BPLY_OT_ToHead)
    bpy.utils.register_class(BPLY_OT_AssetsLoader)
    bpy.utils.register_class(BPLY_OT_AssetsImporter)
    bpy.utils.register_class(BPLY_Preferences)
    bpy.utils.register_class(BPLY_OT_InstallAssets)
    
    bpy.types.WindowManager.poly = bpy.props.PointerProperty(type=BlenderPolyProps)
    bpy.types.WindowManager.poly_model_previews = bpy.props.EnumProperty(items=enum_previews_from_model_previews)

    for category in blender_poly_category_items:
        pcoll = bpy.utils.previews.new ()
        pcoll.previews_previews = ()
        pcoll.previews_previews_dir = ""
        preview_collections [category[0]] = pcoll

def unregister():
    bpy.utils.unregister_class(BPLY_OT_InstallAssets)
    bpy.utils.unregister_class(BPLY_Preferences)
    bpy.utils.unregister_class(BPLY_OT_AssetsImporter)
    bpy.utils.unregister_class(BPLY_OT_AssetsLoader)
    bpy.utils.unregister_class(BPLY_OT_ToHead)
    bpy.utils.unregister_class(BPLY_OT_DirectImport)
    bpy.utils.unregister_class(BPLY_PT_LayoutPanel)
    bpy.utils.unregister_class(BlenderPolyProps)

    for pcoll in preview_collections.values():
        bpy.utils.previews.remove(pcoll)
    preview_collections.clear()
    
    try:
        del bpy.types.WindowManager.poly_model_previews
        del bpy.types.WindowManager.poly
    except:
        pass

if __name__ == "__main__":
    register()