Back to wand and solid keying

chromakey.py

@@ -1,16 +1,51 @@
 import os
 
-import rawpy
-from PIL import Image, ImageOps
+from sklearn.cluster import KMeans
+from wand.image import Image
+import numpy as np
+
+
+def chroma_key(img, object_colors=1, key_coords=(0,0)):
+    hsv = _hsv_points(img)
+    hsv_cart = _hsv_to_cartesian(hsv)
+    labels = _cluster(hsv_cart, object_colors+1, img)
+    mask = _mask(labels, key_coords)
+    with Image.from_array(mask) as mask_img:
+        img.composite_channel('all_channels', mask_img, 'multiply')
+
+
+def _hsv_points(img):
+    hsv_img = img.clone()
+    hsv_img.transform_colorspace('hsv')
+    pixels = list(hsv_img.export_pixels(channel_map='RGB'))  # this actually means HSV
+    return np.array(pixels).reshape(-1, 3)
+
+
+def _hsv_to_cartesian(hsv):
+    h = hsv[:, 0]
+    s = hsv[:, 1]
+    v = hsv[:, 2]
+
+    x = s * np.cos(h * 2 * np.pi / 255)
+    y = s * np.sin(h * 2 * np.pi / 255)
+
+    return np.column_stack((x, y, v))
+
+
+def _cluster(points, num_clusters, img):
+    kmeans = KMeans(n_clusters=num_clusters)
+    kmeans.fit(points)
+    return kmeans.labels_.reshape(img.height, img.width)
+
+
+def _mask(labels, key_coords):
+    key_label = labels[key_coords]
+    return (labels != key_label).astype(np.uint8) * 255
+
 
 path = 'data/RAW/SC1/BR1'
-filename = os.path.join(path, '00570/back.orf')
+filename = os.path.join(path, '41244/back.orf')
 
-raw = rawpy.imread(filename).postprocess()
-rgba = Image.fromarray(raw).convert('RGBA')
-
-h, _, _ = rgba.convert('HSV').split()
-mask = ImageOps.invert(Image.eval(h, lambda x: 255 if 50 < x < 90 else 0))
-
-rgba.putalpha(mask)
-rgba.save('output.png', 'PNG')
+with Image(filename=filename) as img:
+    chroma_key(img)
+    img.save(filename='masked_image.jpg')

sortpending.py

@@ -1,18 +1,15 @@
 import base64
-import io
 import os
 import shutil
 
 from openai import OpenAI
-
-import rawpy
-from PIL import Image
+from wand.image import Image
 
 client = OpenAI()
 
 # List all files and directories in the specified path
 path = 'data/RAW/SC1/BR1'
-pending = os.path.join(path, '00570')
+pending = os.path.join(path, 'pending')
 
 bitting = None
 
@@ -23,13 +20,11 @@ for file in sorted(os.listdir(pending)):
     print(file)
     filename = os.path.join(pending, file)
 
-    raw = rawpy.imread(filename).postprocess()
-    rgb = Image.fromarray(raw).convert('RGB')
-    smaller = rgb.resize((512, 384))
-
-    buf = io.BytesIO()
-    smaller.save(buf, format='JPEG')
-    blob_b64 = base64.b64encode(buf.getvalue()).decode('utf-8')
+    with Image(filename=filename) as img:
+        img.format = 'jpg'
+        img.resize(512, 384)
+        blob = img.make_blob()
+        blob_b64 = base64.b64encode(blob).decode('utf-8')
 
     chat_completion = client.chat.completions.create(
         messages=[