Merge pull request #1 from Oxid15/develop

0.1.0

Author: Oxid15

README.md

@@ -1,2 +1,35 @@
 # ArtworkRetrievalDataGen
+
 Artwork image retrieval synthetic data generator
+
+## Installation
+
+You will need to have `blender` 3.5 installed.  
+  
+Import image as planes addon is used and needs to be activated.  
+  
+To install the project, clone the repo.
+
+```bash
+git clone https://github.com/Oxid15/ArtworkRetrievalDataGen.git
+```
+
+## Usage
+
+1. Create the folder with source images and the one for results
+
+    ```bash
+    mkdir src
+    mkdir dst
+    ```
+
+2. Adjust parameters in the file `generate.py`
+3. Run the rendering
+
+    ```bash
+    <path-to-your-blender>/blender -b -P <path-to-this-repo>/ArtworkRetrievalDataGen/generate.py
+    ```
+
+4. Review the results in the destination folder
+
+    ![Alt text](./examples/00000.png)

config.yaml

@@ -0,0 +1,5 @@
+from .generator import Generator
+
+__version__ = "0.1.0"
+__author__ = "Ilia Moiseev"
+__author_email__ = "ilia.moiseev.5@yandex.ru"

datagen/__init__.py

@@ -0,0 +1,132 @@
+import os
+
+import bpy
+
+from .utils import *
+
+
+class Generator:
+    def __init__(
+        self,
+        src,
+        dest,
+        camera_angles_range=((90, 0), (90, 0)),
+        camera_distance_range=(5, 5),
+        render_per_input=1,
+    ) -> None:
+        self.src = src
+        self.dest = dest
+        self.camera_angles_range = camera_angles_range
+        self.camera_distance_range = camera_distance_range
+        self.render_per_input = render_per_input
+
+        self.parse_images()
+
+    def parse_images(self):
+        self.img_names = os.listdir(self.src)
+
+    def _clear_scene(self):
+        while len(bpy.data.objects):
+            obj = bpy.data.objects[-1]
+            obj.select_set(True)
+            bpy.ops.object.delete(use_global=False)
+
+    def _add_camera(self):
+        ang = self.camera_angles_range
+        first_min, first_max = ang[0][0], ang[1][0]
+        second_min, second_max = ang[0][1], ang[1][1]
+        first_ang = np.random.random() * (first_max - first_min) + first_min
+        second_ang = np.random.random() * (second_max - second_min) + second_min
+
+        d_max, d_min = self.camera_distance_range
+        dist = np.random.random() * (d_max - d_min) + d_min
+
+        loc = spherical2cartesian((deg2rad(first_ang), deg2rad(second_ang)), dist)
+
+        bpy.ops.object.camera_add(
+            align="VIEW",
+            enter_editmode=False,
+            location=loc,
+            rotation=((0.0, 0.0, 0.0)),
+        )
+
+        cam = bpy.data.objects["Camera"]
+        constraint = cam.constraints.new(type="LIMIT_LOCATION")
+        constraint.use_min_z = True
+
+        constraint = cam.constraints.new("TRACK_TO")
+        constraint.target = self.image
+        constraint.track_axis = "TRACK_NEGATIVE_Z"
+        constraint.up_axis = "UP_Y"
+
+    def _add_lights(self):
+        bpy.ops.object.light_add(
+            type="POINT", radius=1, align="VIEW", location=(2, -2, 2)
+        )
+
+    def _add_objects(self, img_name):
+        bpy.ops.import_image.to_plane(
+            files=[{"name": img_name}],
+            directory=self.src,
+            filter_image=True,
+            filter_movie=True,
+            filter_folder=False,
+            relative=False,
+            location=(0, 0, 0),
+        )
+        name = img_name.split(".")[0]
+        self.image = bpy.data.objects[name]
+        self.image.rotation_euler[0] = deg2rad(90)
+        self.image.rotation_euler[2] = deg2rad(90)
+
+        # Wall
+        bpy.ops.mesh.primitive_plane_add(
+            size=1, align="VIEW", enter_editmode=False, location=(-0.01, 0, 0)
+        )
+        plane = bpy.data.objects["Plane"]
+        plane.scale[0] = 1000
+        plane.scale[1] = 1000
+        plane.rotation_euler[0] = deg2rad(90)
+        plane.rotation_euler[2] = deg2rad(90)
+
+        bpy.ops.material.new()
+        bpy.data.materials[-1].node_tree.nodes["Principled BSDF"].inputs[
+            0
+        ].default_value = np.random.random(4)
+        plane.data.materials.append(bpy.data.materials[-1])
+
+        # Floor
+        bpy.ops.mesh.primitive_plane_add(
+            size=1, align="VIEW", enter_editmode=False, location=(0, 0, -0.7)
+        )
+        plane = bpy.data.objects["Plane.001"]
+        plane.scale[0] = 100
+        plane.scale[1] = 100
+
+        bpy.ops.material.new()
+        bpy.data.materials[-1].node_tree.nodes["Principled BSDF"].inputs[
+            0
+        ].default_value = np.random.random(4)
+        plane.data.materials.append(bpy.data.materials[-1])
+
+    def _arrange_scene(self, img_name):
+        self._clear_scene()
+        self._add_lights()
+        self._add_objects(img_name)
+        self._add_camera()
+
+    def _render(self, index, img_name):
+        bpy.context.scene.camera = bpy.data.objects["Camera"]
+
+        base, ext = img_name.split(os.extsep)
+        img_name = base + "_{0:0>5d}".format(index)
+        ".".join((img_name, ext))
+
+        bpy.context.scene.render.filepath = os.path.join(self.dest, img_name)
+        bpy.ops.render.render("INVOKE_DEFAULT", write_still=True)
+
+    def run(self):
+        for name in self.img_names:
+            for i in range(self.render_per_input):
+                self._arrange_scene(name)
+                self._render(i, name)

datagen/generator.py

@@ -1,158 +1,35 @@
 import os
 import sys
-import argparse
-import yaml
-
-import bpy
 
 SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+BASE_DIR = os.path.dirname(SCRIPT_DIR)
 sys.path.append(SCRIPT_DIR)
 
-from utils import *
-
-
-class Generator():
-    def parse_images(self):
-        self.img_names = os.listdir(self.src)
-
-    def _parse_config(self, config_path):
-        with open(config_path, 'r') as f:
-            cfg = yaml.safe_load(f)
-
-        # Required parameters
-        if 'source_folder' in cfg:
-            self.src = cfg['source_folder']
-        else:
-            raise KeyError('"source_folder" is not found in {config_path}')
-
-        if 'destination_folder' in cfg:
-            self.dest = cfg['destination_folder']
-        else:
-            raise KeyError('"destination_folder" is not found in {config_path}')
-
-        # Not required parameters
-        if 'camera_angles_range' in cfg:
-            self.camera_angles_range = [[b for b in a] for a in cfg['camera_angles_range']]
-        else:
-            self.camera_angles_range = ((90, 0), (90, 0))
-        
-        if 'camera_distance_range' in cfg:
-            self.camera_distance_range = cfg['camera_distance_range']
-        else:
-            self.camera_distance_range = (5, 5)
-        
-        if 'render_per_input' in cfg:
-            self.render_per_input = cfg['render_per_input']
-        else:
-            self.render_per_input = 1
-
-        self.parse_images()
-
-    def __init__(self, config_path) -> None:
-        self._parse_config(config_path)
-
-    def _clear_scene(self):
-        while(len(bpy.data.objects)):
-            obj = bpy.data.objects[-1]
-            obj.select = True
-            bpy.ops.object.delete(use_global=False)
-
-    def _add_camera(self):
-        ang = self.camera_angles_range
-        first_min, first_max = ang[0][0], ang[1][0]
-        second_min, second_max = ang[0][1], ang[1][1]
-        first_ang = np.random.random() * (first_max - first_min) + first_min
-        second_ang = np.random.random() * (second_max - second_min) + second_min
-
-        d_max, d_min = self.camera_distance_range
-        dist = np.random.random() * (d_max - d_min) + d_min
-
-        loc = spherical2cartesian((deg2rad(first_ang), deg2rad(second_ang)), dist)
-
-        bpy.ops.object.camera_add(
-            view_align=True,
-            enter_editmode=False,
-            location=loc,
-            rotation=((0., 0., 0.)))
-
-        cam = bpy.data.objects['Camera']
-        constraint = cam.constraints.new(type='LIMIT_LOCATION')
-        constraint.use_min_z = True
-
-        constraint = cam.constraints.new('TRACK_TO')
-        constraint.target = self.image
-        constraint.track_axis = 'TRACK_NEGATIVE_Z'
-        constraint.up_axis = 'UP_Y'
-
-    def _add_lights(self):
-        bpy.ops.object.lamp_add(type='POINT', 
-            radius=1, view_align=False, 
-            location=(2, -2, 2))
-
-    def _add_objects(self, img_name):
-        bpy.ops.import_image.to_plane(
-            files=[{"name": img_name}], 
-            directory=self.src,
-            filter_image=True, filter_movie=True, filter_glob="", relative=False,
-            location=(0,0,0))
-        name = img_name.split('.')[0]
-        self.image = bpy.data.objects[name]
-        self.image.rotation_euler[0] = deg2rad(90)
-        self.image.rotation_euler[2] = deg2rad(90)
-        
-        # Wall
-        bpy.ops.mesh.primitive_plane_add(radius=1, view_align=False, enter_editmode=False, 
-            location=(-0.01, 0, 0))
-        plane = bpy.data.objects['Plane']
-        y_scale = np.random.random() * 3 + 1
-        plane.scale[0] = y_scale
-        plane.rotation_euler[0] = deg2rad(90)
-        plane.rotation_euler[2] = deg2rad(90)
-
-        bpy.ops.material.new()
-        plane.data.materials.append(bpy.data.materials[0])
-        plane.data.materials[0].diffuse_color = np.random.random(3)
-
-        # Floor
-        bpy.ops.mesh.primitive_plane_add(radius=1, view_align=False, enter_editmode=False, 
-            location=(0, 0, -0.7))
-        plane = bpy.data.objects['Plane.001']
-        plane.scale[0] = 10
-        plane.scale[1] = 10
-
-        bpy.ops.material.new()
-        plane.data.materials.append(bpy.data.materials[1])
-        plane.data.materials[0].diffuse_color = np.random.random(3)
-
+from datagen import Generator
 
-    def _arrange_scene(self, img_name):
-        self._clear_scene()
-        self._add_lights()
-        self._add_objects(img_name)
-        self._add_camera()
 
-    def _render(self, index, img_name):
-        bpy.context.scene.camera = bpy.data.objects['Camera']
+# The folder with images to render
+source_folder = os.path.join(BASE_DIR, "src")
 
-        base, ext = img_name.split(os.extsep)
-        img_name = base + "_{0:0>5d}".format(index)
-        '.'.join((img_name, ext))
+# The folder to place renders
+destination_folder = os.path.join(BASE_DIR, "dst")
 
-        bpy.context.scene.render.filepath = os.path.join(self.dest, img_name)
-        bpy.ops.render.render('INVOKE_DEFAULT', write_still=True)
+# The angles and distance are the spherical coordinates of the camera
+# given that the position of the object in the scene is (0,0,0)
+# render_per_input times different angles and distances are sampled
+# from these ranges
+camera_angles_range = [[70, -60], [110, 60]]
 
-    def mainloop(self):
-        for name in self.img_names:
-            for i in range(self.render_per_input):
-                self._arrange_scene(name)
-                self._render(i, name)
+camera_distance_range = [3, 5]
 
+render_per_input = 2
 
-argv = sys.argv[sys.argv.index('--') + 1:]
-parser = argparse.ArgumentParser()
-parser.add_argument('-c', '--cfg', type=str)
-args = parser.parse_known_args(argv)[0]
 
-if __name__ == '__main__':
-    gen = Generator(args.cfg)
-    gen.mainloop()
+if __name__ == "__main__":
+    gen = Generator(
+        src=source_folder,
+        dest=destination_folder,
+        camera_angles_range=camera_angles_range,
+        camera_distance_range=camera_distance_range,
+        render_per_input=render_per_input,
+    ).run()

utils.py datagen/utils.pyutils.py renamed to datagen/utils.py

@@ -1 +1 @@
-pyyaml
+bpy==3.5.0