drsm79 · Sep 15, 2020 · Sep 15, 2020 · Sep 15, 2020 · Sep 16, 2020 · Sep 16, 2020
Showing with 77 additions and 18 deletions.

+9 −0 README.md

BIN _images/modules_668a1a57.jpg

+2 −2 collect.py

+63 −15 randomise.py

+3 −1 requirements.txt
diff --git a/README.md b/README.md
@@ -0,0 +1,9 @@
+# modulair
+
+These are scripts to fetch images of modules from modular sythesizers and format them for training in an object detection model.
+
+![Example training image of synth modules](_images/modules_668a1a57.jpg)
+
+You can read all about this mini-project here:
+
+https://dev.to/hammertoe/using-machine-learning-to-catalog-modular-synthesizers-co2
diff --git a/_images/modules_668a1a57.jpg b/_images/modules_668a1a57.jpg
diff --git a/collect.py b/collect.py
@@ -63,7 +63,7 @@ def save_images(data):
         print(module["id"], module["image"])
         r = requests.get(module["image"], stream=True)
         r.raise_for_status()
-        with open(module["image"].split('/')[-1], 'wb') as f:
+        with open("modules/" + module["image"].split('/')[-1], 'wb') as f:
             for chunk in r:
                 f.write(chunk)
 
@@ -77,7 +77,7 @@ def save_images(data):
         try:
             results = get_search(next)
             data = parse_results(results)
-            with open(f'modules_page_{next}.json', 'w') as f:
+            with open(f'modules/modules_page_{next}.json', 'w') as f:
                 json.dump(data, f)
             save_images(data)
             time.sleep(2)

diff --git a/randomise.py b/randomise.py
@@ -3,9 +3,13 @@
 import hashlib
 import pathlib
 import random
+from tqdm import tqdm
+
 
 from PIL import Image
 
+import tensorflow.compat.v1 as tf
+from object_detection.utils import dataset_util, label_map_util
 
 def join_modules(modules):
     image = False
@@ -36,14 +40,14 @@ def concat_images(height, width, images):
     return dst, coords
 
 
-def pick_modules_from_dir(data_dir='data', count=1):
+def pick_modules_from_dir(data_dir='modules', count=1):
     data = list(pathlib.Path(data_dir).glob('./*.jpg'))
     result = []
-    i = 0
-    while i < count:
+
+    for _ in range(count):
         p = random.choice(data)
         result.append(Image.open(p.as_posix()))
-        i += 1
+
     return result
 
 
@@ -58,29 +62,73 @@ def pick_modules_from_data(jsonfile='data/modules_page_1.json', count=1):
         if p["id"] not in result:
             if '1u' not in p["name"].lower():
                 result[p["id"]] = p
-                image = pathlib.Path(f'data/{p["image"].split("/")[-1]}')
+                image = pathlib.Path(f'modules/{p["image"].split("/")[-1]}')
                 if image.exists():
                     result[p["id"]]["image"] = Image.open(image.as_posix())
                     i += 1
     return result
 
 
-if __name__ == "__main__":
-    modules = pick_modules_from_data(count=5)
+def gen_tfrecord():
+    page = random.randint(1, 5)
+    jsonfile=f'modules/modules_page_{page}.json'
+    modules = pick_modules_from_data(count=5, jsonfile=jsonfile)
     image, all_coords = join_modules([m["image"] for m in modules.values()])
-
+    
     hash = hashlib.sha256()
     for i in [m.encode('utf-8') for m in modules]:
         hash.update(bytes(i))
     h = hash.hexdigest()[:8]
-
-    filename = f'modules_{h}.jpg'
+    filename = f'composites/modules_{h}.jpg'
     image.save(filename)
-    to_save = []
+
+    xmins, xmaxs = [], []
+    ymins, ymaxs = [], []
+    classes_text, classes = [], []
 
     for i, d in enumerate(modules.values()):
         data = {k: v for k, v in d.items() if k != 'image'}
-        data['x_min'], data['x_max'] = all_coords[i]
-        to_save.append(data)
-    with open(f'modules_{h}.json', 'w') as f:
-        json.dump(to_save, f, indent=2, sort_keys=True)
+        x_min, x_max = all_coords[i]
+
+        xmins.append(x_min)
+        xmaxs.append(x_max)
+        ymins.append(0)
+        ymaxs.append(image.height)
+        classes_text.append(data['name'].encode('utf-8'))
+        classes.append(int(data['id']))
+
+    tf_example = tf.train.Example(features=tf.train.Features(feature={
+        'image/height': dataset_util.int64_feature(image.height),
+        'image/width': dataset_util.int64_feature(image.width),
+        'image/filename': dataset_util.bytes_feature(filename.encode('utf-8')),
+        'image/source_id': dataset_util.bytes_feature(filename.encode('utf-8')),
+        'image/encoded': dataset_util.bytes_feature(open(filename, "rb").read()),
+        'image/format': dataset_util.bytes_feature( b'jpg'),
+        'image/object/bbox/xmin': dataset_util.float_list_feature(xmins),
+        'image/object/bbox/xmax': dataset_util.float_list_feature(xmaxs),
+        'image/object/bbox/ymin': dataset_util.float_list_feature(ymins),
+        'image/object/bbox/ymax': dataset_util.float_list_feature(ymaxs),
+        'image/object/class/text': dataset_util.bytes_list_feature(classes_text),
+        'image/object/class/label': dataset_util.int64_list_feature(classes),
+    }))
+
+    return tf_example
+
+
+def write_records(filename, num):
+
+    writer = tf.python_io.TFRecordWriter(filename)
+
+    for i in tqdm(range(num)):
+        tf_example = gen_tfrecord()
+        writer.write(tf_example.SerializeToString())
+    writer.close()
+
+    print(f'Successfully created the TFRecord file with {num} records: {filename}')
+
+
+if __name__ == "__main__":
+
+    write_records("train.record", 1000)
+    write_records("test.record", 100)
+
diff --git a/requirements.txt b/requirements.txt
@@ -1,4 +1,6 @@
 beautifulsoup4==4.9.1
 Pillow==7.2.0
 requests==2.24.0
-tensorflow
+tensorflow
+tqdm
+