From 35d78bf37e73b59aac01055d48cb2ffab36c13e8 Mon Sep 17 00:00:00 2001
From: Ian Gulliver <github@flamingcow.io>
Date: Sun, 11 Aug 2024 14:28:00 -0700
Subject: [PATCH] Switch to rawpy/PIL and simplify

---
 chromakey.py   | 19 ++++------
 sortpending.py | 97 ++++++++++++++++++++++++++------------------------
 2 files changed, 57 insertions(+), 59 deletions(-)

diff --git a/chromakey.py b/chromakey.py
index df42350..c856d4f 100644
--- a/chromakey.py
+++ b/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)
-rgb_image = raw.postprocess()
+h, _, _ = rgba.convert('HSV').split()
+mask = ImageOps.invert(Image.eval(h, lambda x: 255 if 50 < x < 90 else 0))
 
-image = Image.fromarray(rgb_image).convert('RGBA')
-
-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')
\ No newline at end of file
+rgba.putalpha(mask)
+rgba.save('output.png', 'PNG')
\ No newline at end of file
diff --git a/sortpending.py b/sortpending.py
index e050346..c91a088 100644
--- a/sortpending.py
+++ b/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:
-        img.format = 'jpg'
-        img.resize(512, 384)
-        blob = img.make_blob()
-        blob_b64 = base64.b64encode(blob).decode('utf-8')
+    raw = rawpy.imread(filename).postprocess()
+    rgb = Image.fromarray(raw).convert('RGB')
+    smaller = rgb.resize((512, 384))
 
-        chat_completion = client.chat.completions.create(
-            messages=[
-                {
-                    'role': 'user',
-                    'content': [
-                        {
-                            '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.'
-                        },
-                        {
-                            'type': 'image_url',
-                            'image_url': {
-                                'url': f'data:image/jpeg;base64,{blob_b64}'
-                            }
+    buf = io.BytesIO()
+    smaller.save(buf, format='JPEG')
+    blob_b64 = base64.b64encode(buf.getvalue()).decode('utf-8')
+
+    chat_completion = client.chat.completions.create(
+        messages=[
+            {
+                'role': 'user',
+                'content': [
+                    {
+                        '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.'
+                    },
+                    {
+                        '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:
-            case 'front':
-                assert bitting is not None, 'key image before card image'
-                dest = os.path.join(path, bitting, 'front.orf')
-                assert not os.path.exists(dest), 'key image already exists'
-                shutil.move(filename, dest)
+    match content:
+        case 'front':
+            assert bitting is not None, 'key image before card image'
+            dest = os.path.join(path, bitting, 'front.orf')
+            assert not os.path.exists(dest), 'key image already exists'
+            shutil.move(filename, dest)
 
-            case 'back':
-                assert bitting is not None, 'key image before card image'
-                dest = os.path.join(path, bitting, 'back.orf')
-                assert not os.path.exists(dest), 'key image already exists'
-                shutil.move(filename, dest)
+        case 'back':
+            assert bitting is not None, 'key image before card image'
+            dest = os.path.join(path, bitting, 'back.orf')
+            assert not os.path.exists(dest), 'key image already exists'
+            shutil.move(filename, dest)
 
-            case _:
-                bitting = content
-                os.makedirs(os.path.join(path, bitting), exist_ok=True)
-                dest = os.path.join(path, bitting, 'card.orf')
-                assert not os.path.exists(dest), 'card image already exists'
-                shutil.move(filename, dest)
\ No newline at end of file
+        case _:
+            bitting = content
+            os.makedirs(os.path.join(path, bitting), exist_ok=True)
+            dest = os.path.join(path, bitting, 'card.orf')
+            assert not os.path.exists(dest), 'card image already exists'
+            shutil.move(filename, dest)
\ No newline at end of file