File size: 2,933 Bytes
4373daa
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
import gradio as gr
import bpy
from tqdm import tqdm
from math import pi
import tempfile
import molecularnodes as mn
import os


window = bpy.context.window
screen = window.screen

style = 'cartoon'

nodes_to_append = ["MN_color_set", 
                   "MN_color_common", 
                   "MN_color_attribute_random",
                   mn.nodes.styles_mapping[style]]

def get_areas(type):
    return [area for area in screen.areas if area.type == type]

def get_regions(areas):
    return [region for region in areas[0].regions if region.type == 'WINDOW']

for node in nodes_to_append:
    bpy.ops.wm.append(
                'INVOKE_DEFAULT',
                directory = os.path.join(mn.nodes.mn_data_file, 'NodeTree'), 
                filename = node, 
                link = False
            )

def generate(progress=gr.Progress(track_tqdm=True)):
    area_type = 'VIEW_3D' 
    areas  = get_areas(area_type)
    with bpy.context.temp_override(window=window, area=areas[0], region=get_regions(areas)[0], screen=screen):
        
        
        for obj in bpy.context.scene.objects:
            if obj.type == 'MESH':
                bpy.data.objects.remove(obj, do_unlink=True)

        molecule = mn.load.molecule_rcsb("7TYG", starting_style=style, center_molecule=True) # <- this line would causes the error 🐛
        molecule.select_set(True)
        bpy.context.view_layer.objects.active = molecule

        bpy.ops.view3d.camera_to_view_selected()
        camera = bpy.data.objects["Camera"]
        camera.data.dof.use_dof = True
        camera.data.dof.focus_distance = 5
        camera.data.dof.aperture_fstop = 4
        camera.data.angle = pi / 3
        camera.data.type = "PERSP"

        with tempfile.NamedTemporaryFile(suffix=".JPEG", delete=False) as f:
            bpy.context.scene.render.resolution_y = 1000
            bpy.context.scene.render.resolution_x = 1000
            bpy.context.scene.render.image_settings.file_format = "JPEG"
            bpy.context.scene.render.filepath = f.name

            with tqdm() as pbar:

                def elapsed(dummy):
                    pbar.update()

                bpy.app.handlers.render_stats.append(elapsed)
                bpy.context.scene.frame_set(1)
                bpy.context.scene.frame_current = 1
                bpy.ops.render.render(animation=False, write_still=True)

                bpy.data.images["Render Result"].save_render(
                    filepath=bpy.context.scene.render.filepath
                )
                bpy.app.handlers.render_stats.clear()
                return f.name

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            render_btn = gr.Button("Render")
        with gr.Column(scale=3):
            image = gr.Image(type="filepath")

    render_btn.click(
        generate,
        outputs=[image],
    )

demo.queue(concurrency_count=1)
demo.launch(debug=True, inline=True)