diff --git a/chromakey.py b/chromakey.py index c856d4f..40f6fa0 100644 --- a/chromakey.py +++ b/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') \ No newline at end of file +with Image(filename=filename) as img: + chroma_key(img) + img.save(filename='masked_image.jpg') \ No newline at end of file diff --git a/sortpending.py b/sortpending.py index c91a088..b2e6140 100644 --- a/sortpending.py +++ b/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=[