meshrcnn_overlay_mesh.py

# Blender 2.93 python script for producing images like in https://github.com/facebookresearch/meshrcnn/issues/22
# blender -noaudio --background --python meshrcnn_overlay_mesh.py -- --object-path test.obj --background-path test.png
#
# Sample ouptut:
# https://github.com/facebookresearch/meshrcnn/issues/100
#
# References:
# https://github.com/yuki-koyama/blender-cli-rendering/blob/master/02_suzanne.py
# https://github.com/weiaicunzai/blender_shapenet_render/blob/master/render_rgb.py
# https://towardsdatascience.com/blender-2-8-grease-pencil-scripting-and-generative-art-cbbfd3967590
# https://github.com/5agado/data-science-learning/tree/master/graphics
# https://github.com/mikedh/trimesh/blob/3b0c34df50142ce0e54f20113164fca818c9696f/trimesh/bounds.py
# https://b3d.interplanety.org/en/render-from-console-only-on-the-specified-gpu-devices-2/
#
# Demo:
# git clone https://github.com/facebookresearch/meshrcnn
# cd meshrcnn
# pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu111/torch1.9/index.html
# pip install pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/py39_cu111_pyt190/download.html
# pip install -e .
# wget -O test.jpg https://user-images.githubusercontent.com/4369065/77708628-cda99d00-6f85-11ea-949a-5dad891005ee.jpg
# python demo/demo.py --input test.jpg --config-file configs/pix3d/meshrcnn_R50_FPN.yaml --output output_demo --onlyhighest MODEL.WEIGHTS meshrcnn://meshrcnn_R50.pth

import os
import sys
import math
import argparse
import subprocess
try:
    import numpy as np, scipy.spatial 
except:
    # install dependencies in Blender python
    subprocess.call([sys.executable, '-m', 'ensurepip'])
    subprocess.call([sys.executable, '-m', 'pip', 'install', 'numpy', 'scipy'])
finally:
    import numpy as np, scipy.spatial

import bpy

def unitize(vectors, threshold = np.finfo(np.float64).resolution * 100):
    vectors = np.asanyarray(vectors)

    # for (m, d) arrays take the per-row unit vector
    # using sqrt and avoiding exponents is slightly faster
    # also dot with ones is faser than .sum(axis=1)
    norm = np.sqrt(np.dot(vectors * vectors, [1.0] * vectors.shape[1]))
    valid = norm > threshold
    norm[valid] **= -1
    unit = vectors * norm.reshape((-1, 1))

    return unit[valid]

def planar_matrix(offset=[0.0, 0.0], theta=0.0):
    """
    2D homogeonous transformation matrix.
    Parameters
    ----------
    offset : (2,) float
      XY offset
    theta : float
      Rotation around Z in radians
    Returns
    ----------
    matrix : (3, 3) flat
      Homogeneous 2D transformation matrix
    """
    offset = np.asanyarray(offset, dtype=np.float64)
    theta = float(theta)
    
    T = np.eye(3, dtype=np.float64)
    s = np.sin(theta)
    c = np.cos(theta)

    T[0, :2] = [c, s]
    T[1, :2] = [-s, c]
    T[:2, 2] = offset

    return T


def oriented_bounds_2D(points, qhull_options='QbB'):
    """
    Find an oriented bounding box for an array of 2D points.
    Parameters
    ----------
    points : (n,2) float
      Points in 2D.
    Returns
    ----------
    transform : (3,3) float
      Homogeneous 2D transformation matrix to move the
      input points so that the axis aligned bounding box
      is CENTERED AT THE ORIGIN.
    rectangle : (2,) float
       Size of extents once input points are transformed
       by transform
    """
    # make sure input is a numpy array
    points = np.asanyarray(points, dtype=np.float64)
    # create a convex hull object of our points
    # 'QbB' is a qhull option which has it scale the input to unit
    # box to avoid precision issues with very large/small meshes
    convex = scipy.spatial.ConvexHull(points, qhull_options=qhull_options)

    # (n,2,3) line segments
    hull_edges = convex.points[convex.simplices]
    # (n,2) points on the convex hull
    hull_points = convex.points[convex.vertices]

    # unit vector direction of the edges of the hull polygon
    # filter out zero- magnitude edges via check_valid
    edge_vectors = unitize(np.diff(hull_edges, axis=1).reshape((-1, 2)))

    # create a set of perpendicular vectors
    perp_vectors = np.fliplr(edge_vectors) * [-1.0, 1.0]

    # find the projection of every hull point on every edge vector
    # this does create a potentially gigantic n^2 array in memory,
    # and there is the 'rotating calipers' algorithm which avoids this
    # however, we have reduced n with a convex hull and numpy dot products
    # are extremely fast so in practice this usually ends up being pretty
    # reasonable
    x = np.dot(edge_vectors, hull_points.T)
    y = np.dot(perp_vectors, hull_points.T)

    # reduce the projections to maximum and minimum per edge vector
    bounds = np.column_stack((x.min(axis=1),
                              y.min(axis=1),
                              x.max(axis=1),
                              y.max(axis=1)))

    # calculate the extents and area for each edge vector pair
    extents = np.diff(bounds.reshape((-1, 2, 2)),
                      axis=1).reshape((-1, 2))
    area = np.product(extents, axis=1)
    area_min = area.argmin()

    # (2,) float of smallest rectangle size
    rectangle = extents[area_min]

    # find the (3,3) homogeneous transformation which moves the input
    # points to have a bounding box centered at the origin
    offset = -bounds[area_min][:2] - (rectangle * .5)
    theta = np.arctan2(*edge_vectors[area_min][::-1])
    transform = planar_matrix(offset, theta)

    # we would like to consistently return an OBB with
    # the largest dimension along the X axis rather than
    # the long axis being arbitrarily X or Y.
    if np.less(*rectangle):
        # a 90 degree rotation
        flip = planar_matrix(theta=np.pi / 2)
        # apply the rotation
        transform = np.dot(flip, transform)
        # switch X and Y in the OBB extents
        rectangle = np.roll(rectangle, 1)

    w, h = rectangle
    points = np.array([(-w/2.0, -h/2.0, 1.0), (-w/2.0, h/2.0, 1.0), (w/2.0, -h/2.0, 1.0), (w/2.0, h/2.0, 1.0)], dtype = np.float64)

    points = np.linalg.inv(transform) @ points.T
    return points[:2].T.tolist()
  
def configure_cycles_renderer(scene, camera_object, num_samples = 16, use_denoising = True, use_motion_blur = False, use_transparent_bg = False, prefer_cuda_use = True, use_adaptive_sampling = False):
    scene.render.image_settings.file_format = 'PNG'
    scene.render.engine = 'CYCLES'
    scene.render.use_motion_blur = use_motion_blur
    scene.render.film_transparent = use_transparent_bg
    scene.cycles.use_adaptive_sampling = use_adaptive_sampling
    scene.cycles.samples = num_samples
    scene.view_layers[0].cycles.use_denoising = use_denoising

    # Source - https://blender.stackexchange.com/a/196702
    if prefer_cuda_use:
        bpy.context.scene.cycles.device = 'GPU'
        bpy.context.preferences.addons['cycles'].preferences.compute_device_type = 'CUDA'
    bpy.context.preferences.addons['cycles'].preferences.get_devices()
    for d in bpy.context.preferences.addons['cycles'].preferences.devices:
        d.use = False
        print('-', d.name, d.use)
        
def configure_compositing(background_path, output_path):
    bpy.context.scene.use_nodes = True
    tree = bpy.context.scene.node_tree
    links = tree.links

    for node in tree.nodes:
        tree.nodes.remove(node)
    
    image_node = tree.nodes.new('CompositorNodeImage')
    scale_node = tree.nodes.new('CompositorNodeScale')
    alpha_over_node = tree.nodes.new('CompositorNodeAlphaOver')
    render_layer_node = tree.nodes.new('CompositorNodeRLayers')
    file_output_node = tree.nodes.new('CompositorNodeOutputFile')

    #enum in [‘RELATIVE’, ‘ABSOLUTE’, ‘SCENE_SIZE’, ‘RENDER_SIZE’], default ‘RELATIVE’
    scale_node.space = 'RENDER_SIZE'
    image_node.image = bpy.data.images.load(background_path)
    file_output_node.base_path = os.path.dirname(output_path)
    file_output_node.file_slots[0].path = os.path.basename(output_path)

    links.new(image_node.outputs[0], scale_node.inputs[0])
    links.new(scale_node.outputs[0], alpha_over_node.inputs[1])
    links.new(render_layer_node.outputs[0], alpha_over_node.inputs[2])
    links.new(alpha_over_node.outputs[0], file_output_node.inputs[0])

def draw_line(gp_frame, p0, p1, line_width = 1000):
    gp_stroke = gp_frame.strokes.new()
    gp_stroke.line_width = line_width
    gp_stroke.display_mode = '3DSPACE'  # allows for editing
    gp_stroke.points.add(count=2)
    gp_stroke.points[0].co = p0
    gp_stroke.points[1].co = p1
        

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--object-path', default = '/Users/kanvadim/Downloads/6fcaaaaf9f19895ffe8b20c86c900a99/test.obj')
    parser.add_argument('--background-path', default = '/Users/kanvadim/Downloads/6fcaaaaf9f19895ffe8b20c86c900a99/test.png')
    parser.add_argument('--lens', type = float, default = 20)
    parser.add_argument('--cuda', action = 'store_true')
    parser.add_argument('--box3d-requiring-display', action = 'store_true', help = 'https://developer.blender.org/T75952')
    args = parser.parse_args(sys.argv[1 + sys.argv.index('--'):] if '--' in sys.argv else [])
    
    for item in bpy.data.objects:
        bpy.data.objects.remove(item)

    bpy.ops.import_scene.obj(filepath = args.object_path)
    for obj in bpy.context.scene.objects:
        if obj.type == 'MESH':
            verts = [vert.co.to_tuple() for vert in obj.data.vertices]
            # do we need to reset rotation?
            obj.rotation_euler = (0, 0, 0)
            break
    
    verts_proj = [(v[0], v[-1]) for v in verts]
    miny, maxy = min(v[1] for v in verts), max(v[1] for v in verts)
    bbox_points_2D = oriented_bounds_2D(verts_proj)

    verts = [(v[0], miny, v[-1]) for v in bbox_points_2D] + [(v[0], maxy, v[-1]) for v in bbox_points_2D]
    
    edges = [
        (0, 1), 
        (1, 3), 
        (3, 2),
        (2, 0),
        (4, 5),
        (5, 7),
        (7, 6),
        (6, 4),
        (0, 4),
        (1, 5),
        (2, 6),
        (3, 7)
    ]
    
    if args.box3d_requiring_display:
        material = bpy.data.materials.new("mymaterial")
        bpy.data.materials.create_gpencil_data(material)
        material.grease_pencil.color = (1, 0, 0, 1)

        #bpy.ops.object.gpencil_add(location=(0, 0, 0), rotation = (math.pi / 2, 0, 0), type='EMPTY')
        bpy.ops.object.gpencil_add(location=(0, 0, 0), rotation = (0, 0, 0), type='EMPTY')
        gp = bpy.context.scene.objects[-1]
        gp.active_material = material
        gp_layer = gp.data.layers.new('gplayer', set_active=True)
        gp_layer.use_lights = False
        gp_frame = gp_layer.frames.new(0)
        for u, v in edges:
            draw_line(gp_frame, verts[u], verts[v])

    bpy.ops.object.light_add(type='SUN', location=(0, 0, 0), rotation=(0, math.pi, 0))
    #bpy.ops.object.light_add(type='SUN', location=(0, 0, 0), rotation=(0, 0, 0))
    
    bpy.ops.object.camera_add(location=(0, 0, 0), rotation=(0, math.pi, 0))
    #bpy.ops.object.camera_add(location=(0, 0, 0), rotation=(math.pi / 2, 0, math.pi))
    bpy.data.cameras[-1].lens = args.lens
    camera = bpy.context.object
    
    res_x, res_y = bpy.data.images.load(args.background_path).size
    
    scene = bpy.data.scenes["Scene"]
    scene.render.resolution_percentage = 100
    scene.render.resolution_x = res_x
    scene.render.resolution_y = res_y
    scene.camera = camera
    
    configure_cycles_renderer(scene, camera, use_transparent_bg = True, prefer_cuda_use = args.cuda)
    
    configure_compositing(args.background_path, args.object_path + '-overlay.png')
    bpy.ops.render.render(write_still = False)