Update app.py

app.py

@@ -4,21 +4,20 @@ import io, base64, struct
 
 app = Flask(__name__, static_folder="static")
 
-WIDTH, HEIGHT = 1184, 1184
+WIDTH, HEIGHT = 1672, 1672
 BG_COLOR = (88, 28, 135)
 TEXT_COLOR = (255, 255, 255)
 BITS_PER_CHANNEL = 2
 MASK = (1 << BITS_PER_CHANNEL) - 1
-CHANNELS = 3
-BITS_PER_PIXEL = BITS_PER_CHANNEL * CHANNELS
+MAX_BYTES = 2 * 1024 * 1024
 
 def make_base_image():
     img = Image.new("RGB", (WIDTH, HEIGHT), BG_COLOR)
     draw = ImageDraw.Draw(img)
     try:
-        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 80)
+        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 200)
     except:
-        font = ImageFont.load_default()
+        font = ImageFont.load_default(size=200)
     text = "CC"
     bbox = draw.textbbox((0, 0), text, font=font)
     tw = bbox[2] - bbox[0]
@@ -28,22 +27,12 @@ def make_base_image():
     draw.text((x, y), text, fill=TEXT_COLOR, font=font)
     return img
 
-def get_text_pixel_set():
-    img = make_base_image()
-    pixels = img.load()
-    text_pixels = set()
-    for y in range(HEIGHT):
-        for x in range(WIDTH):
-            if pixels[x, y] != BG_COLOR:
-                text_pixels.add((x, y))
-    return text_pixels
-
-def get_carrier_pixels():
-    text_pixels = get_text_pixel_set()
+def get_carriers(base_img):
+    pixels = base_img.load()
     carriers = []
     for y in range(HEIGHT):
         for x in range(WIDTH):
-            if (x, y) not in text_pixels:
+            if pixels[x, y] == BG_COLOR:
                 carriers.append((x, y))
     return carriers
 
@@ -66,17 +55,15 @@ def bytes_from_bits(bits):
         result.append(byte)
     return bytes(result)
 
-def encode(text):
-    img = make_base_image()
-    carriers = get_carrier_pixels()
-    payload = text.encode("utf-8")
+def embed(payload):
+    base_img = make_base_image()
+    carriers = get_carriers(base_img)
     header = struct.pack(">I", len(payload))
     data = header + payload
     bits = bits_from_bytes(data)
-    max_bits = len(carriers) * BITS_PER_PIXEL
-    if len(bits) > max_bits:
-        raise ValueError(f"Text too long. Max ~{(max_bits // 8) - 4} bytes.")
-    pixels = img.load()
+    if len(bits) > len(carriers) * BITS_PER_CHANNEL * 3:
+        raise ValueError(f"Data too large. Max {MAX_BYTES} bytes.")
+    pixels = base_img.load()
     bit_idx = 0
     for (x, y) in carriers:
         if bit_idx >= len(bits):
@@ -93,10 +80,11 @@ def encode(text):
                     chunk = chunk << 1
             new_channels.append((c & ~MASK) | chunk)
         pixels[x, y] = tuple(new_channels)
-    return img
+    return base_img
 
-def decode(img):
-    carriers = get_carrier_pixels()
+def extract(img):
+    base_img = make_base_image()
+    carriers = get_carriers(base_img)
     pixels = img.load()
     bits = []
     for (x, y) in carriers:
@@ -104,25 +92,43 @@ def decode(img):
         for c in [r, g, b]:
             for i in range(BITS_PER_CHANNEL - 1, -1, -1):
                 bits.append((c >> i) & 1)
-    header_bits = bits[:32]
-    length = struct.unpack(">I", bytes_from_bits(header_bits))[0]
-    if length == 0 or length > 1_100_000:
-        return None
-    payload_bits = bits[32:32 + length * 8]
-    return bytes_from_bits(payload_bits).decode("utf-8", errors="replace")
+    length = struct.unpack(">I", bytes_from_bits(bits[:32]))[0]
+    if length == 0 or length > MAX_BYTES + 100:
+        return None, None
+    payload = bytes_from_bits(bits[32:32 + length * 8])
+    return length, payload
 
 @app.route("/")
 def index():
     return send_from_directory("static", "index.html")
 
-@app.route("/encode", methods=["POST"])
-def encode_route():
+@app.route("/encode/text", methods=["POST"])
+def encode_text():
     data = request.json
     text = data.get("text", "")
     if not text:
         return jsonify({"error": "No text provided"}), 400
     try:
-        img = encode(text)
+        payload = b'\x00' + text.encode("utf-8")
+        img = embed(payload)
+        buf = io.BytesIO()
+        img.save(buf, format="PNG")
+        img_b64 = base64.b64encode(buf.getvalue()).decode()
+        return jsonify({"image": img_b64})
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+@app.route("/encode/image", methods=["POST"])
+def encode_image():
+    if "file" not in request.files:
+        return jsonify({"error": "No file uploaded"}), 400
+    f = request.files["file"]
+    raw = f.read()
+    if len(raw) > MAX_BYTES:
+        return jsonify({"error": f"Image too large. Max 2MB, got {len(raw) / 1024 / 1024:.2f}MB"}), 400
+    try:
+        payload = b'\x01' + raw
+        img = embed(payload)
         buf = io.BytesIO()
         img.save(buf, format="PNG")
         img_b64 = base64.b64encode(buf.getvalue()).decode()
@@ -136,10 +142,18 @@ def decode_route():
         return jsonify({"error": "No file uploaded"}), 400
     f = request.files["file"]
     img = Image.open(f).convert("RGB")
-    text = decode(img)
-    if text is None:
+    if img.size != (WIDTH, HEIGHT):
+        return jsonify({"error": "Not a valid ChromaCode image"}), 400
+    length, payload = extract(img)
+    if payload is None:
         return jsonify({"error": "Nothing encoded or invalid image"}), 400
-    return jsonify({"text": text})
+    mode = payload[0]
+    content = payload[1:]
+    if mode == 0:
+        return jsonify({"type": "text", "text": content.decode("utf-8", errors="replace")})
+    else:
+        img_b64 = base64.b64encode(content).decode()
+        return jsonify({"type": "image", "image": img_b64})
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)