Switch to rawpy/PIL and simplify

chromakey.py

@@ -6,18 +6,11 @@ from PIL import Image, ImageOps
 path = 'data/RAW/SC1/BR1'
 filename = os.path.join(path, '00570/back.orf')
 
+raw = rawpy.imread(filename).postprocess()
+rgba = Image.fromarray(raw).convert('RGBA')
 
-raw = rawpy.imread(filename)
+h, _, _ = rgba.convert('HSV').split()
-rgb_image = raw.postprocess()
+mask = ImageOps.invert(Image.eval(h, lambda x: 255 if 50 < x < 90 else 0))
 
-image = Image.fromarray(rgb_image).convert('RGBA')
+rgba.putalpha(mask)
-
+rgba.save('output.png', 'PNG')
-hsv_image = image.convert('HSV')
-h, s, v = hsv_image.split()
-
-green_mask = Image.eval(h, lambda x: 255 if 50 < x < 90 else 0)
-green_mask = ImageOps.invert(green_mask)
-
-rgba_image = image.copy()
-rgba_image.putalpha(green_mask)
-rgba_image.save('output.png', 'PNG')

sortpending.py

@@ -1,15 +1,18 @@
-import os
 import base64
+import io
+import os
 import shutil
 
 from openai import OpenAI
-from wand.image import Image
+
+import rawpy
+from PIL import Image
 
 client = OpenAI()
 
 # List all files and directories in the specified path
 path = 'data/RAW/SC1/BR1'
-pending = os.path.join(path, 'pending')
+pending = os.path.join(path, '00570')
 
 bitting = None
 
@@ -20,53 +23,55 @@ for file in sorted(os.listdir(pending)):
     print(file)
     filename = os.path.join(pending, file)
 
-    with Image(filename=filename) as img:
+    raw = rawpy.imread(filename).postprocess()
-        img.format = 'jpg'
+    rgb = Image.fromarray(raw).convert('RGB')
-        img.resize(512, 384)
+    smaller = rgb.resize((512, 384))
-        blob = img.make_blob()
-        blob_b64 = base64.b64encode(blob).decode('utf-8')
 
-        chat_completion = client.chat.completions.create(
+    buf = io.BytesIO()
-            messages=[
+    smaller.save(buf, format='JPEG')
-                {
+    blob_b64 = base64.b64encode(buf.getvalue()).decode('utf-8')
-                    'role': 'user',
+
-                    'content': [
+    chat_completion = client.chat.completions.create(
-                        {
+        messages=[
-                            'type': 'text',
+            {
-                            'text': 'If this is a picture of a white card with black lettering on it, respond with just the 5 digit number written on the card. If this is a picture of a key, respond with "front" if you can see the engravings "SC1" & "USA", and "back" if you cannot.'
+                'role': 'user',
-                        },
+                'content': [
-                        {
+                    {
-                            'type': 'image_url',
+                        'type': 'text',
-                            'image_url': {
+                        'text': 'If this is a picture of a white card with black lettering on it, respond with just the 5 digit number written on the card. If this is a picture of a key, respond with "front" if you can see the engravings "SC1" & "USA", and "back" if you cannot.'
-                                'url': f'data:image/jpeg;base64,{blob_b64}'
+                    },
-                            }
+                    {
+                        'type': 'image_url',
+                        'image_url': {
+                            'url': f'data:image/jpeg;base64,{blob_b64}'
                         }
-                    ]
+                    }
-                }
+                ]
-            ],
+            }
-            model='gpt-4o',
+        ],
-        )
+        model='gpt-4o',
+    )
 
-        content = chat_completion.choices[0].message.content
+    content = chat_completion.choices[0].message.content
 
-        print(content)
+    print(content)
 
-        match content:
+    match content:
-            case 'front':
+        case 'front':
-                assert bitting is not None, 'key image before card image'
+            assert bitting is not None, 'key image before card image'
-                dest = os.path.join(path, bitting, 'front.orf')
+            dest = os.path.join(path, bitting, 'front.orf')
-                assert not os.path.exists(dest), 'key image already exists'
+            assert not os.path.exists(dest), 'key image already exists'
-                shutil.move(filename, dest)
+            shutil.move(filename, dest)
 
-            case 'back':
+        case 'back':
-                assert bitting is not None, 'key image before card image'
+            assert bitting is not None, 'key image before card image'
-                dest = os.path.join(path, bitting, 'back.orf')
+            dest = os.path.join(path, bitting, 'back.orf')
-                assert not os.path.exists(dest), 'key image already exists'
+            assert not os.path.exists(dest), 'key image already exists'
-                shutil.move(filename, dest)
+            shutil.move(filename, dest)
 
-            case _:
+        case _:
-                bitting = content
+            bitting = content
-                os.makedirs(os.path.join(path, bitting), exist_ok=True)
+            os.makedirs(os.path.join(path, bitting), exist_ok=True)
-                dest = os.path.join(path, bitting, 'card.orf')
+            dest = os.path.join(path, bitting, 'card.orf')
-                assert not os.path.exists(dest), 'card image already exists'
+            assert not os.path.exists(dest), 'card image already exists'
-                shutil.move(filename, dest)
+            shutil.move(filename, dest)