From 33d475a1eb6e52f30f2b551a590b234fd982a718 Mon Sep 17 00:00:00 2001
From: jaseg <git@jaseg.de>
Date: Mon, 21 Oct 2024 16:47:21 +0200
Subject: Add P200 model support

---
 example.py         |  2 +-
 src/infiray_irg.py | 45 +++++++++++++++++++++++++++++++++++----------
 2 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/example.py b/example.py
index f434b0b..0cb97cc 100644
--- a/example.py
+++ b/example.py
@@ -3,7 +3,7 @@ from pathlib import Path
 
 import infiray_irg
 
-coarse, fine, vis = infiray_irg.load(Path('example-2.irg').read_bytes())
+coarse, fine, vis = infiray_irg.load(Path('example.irg').read_bytes())
 
 fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 18))
 
diff --git a/src/infiray_irg.py b/src/infiray_irg.py
index 80d04b3..22d500e 100644
--- a/src/infiray_irg.py
+++ b/src/infiray_irg.py
@@ -10,13 +10,17 @@ def load(data):
         if len(out) < n:
             raise ValueError(f'File is truncated, tried to read {n} bytes, but only {len(out)} bytes remain.')
         return out
-        
+
     header = consume(128)
-    if header[:2] not in [
-            bytes([0xca, 0xac]),
-            bytes([0xba, 0xab])
-            ] or header[-2:] != bytes([0xac, 0xca]):
-        raise ValueError('Header magic not found.')
+    for model, match in {
+            'c201': bytes([0xca, 0xac]),
+            'other': bytes([0xba, 0xab]),
+            'p200': bytes([0x04, 0xa0]),
+            }.items():
+        if (header[:2] + header[-2:]).startswith(match):
+            break
+    else:
+        raise ValueError(f'Header magic not found. Got header: {header[0]:02x} {header[1]:02x}')
 
     _unk0, coarse_section_length, y_res, x_res,\
     _zero0, _unk1, _zero1, fine_offset, _unk2, jpeg_length,\
@@ -24,26 +28,47 @@ def load(data):
 
     _zero_celsius0, _zero_celsius1, *rest, high_gain_mode_flag = struct.unpack('<11I', header[34:78])
     
-    if (x_res, y_res) != (x_res_2, y_res_2):
+    import textwrap
+    print(textwrap.dedent(f'''
+        {_unk0=}, {coarse_section_length=}, {y_res=}, {x_res=},
+        {_zero0=}, {_unk1=}, {_zero1=}, {fine_offset=}, {_unk2=}, {jpeg_length=},
+        {y_res_2=}, {x_res_2=}, {_unk3=}
+        {_zero_celsius0=} {_zero_celsius1=} {rest=}, {high_gain_mode_flag=}'''))
+
+    if (x_res, y_res) != (x_res_2, y_res_2) and model != 'p200':
         raise ValueError(f'Resolution mismatch in header: {x_res}*{y_res} != {x_res_2}*{y_res_2}')
 
     if x_res*y_res != coarse_section_length:
         raise ValueError('Resolution mismatch in header')
         
-    coarse_img = np.frombuffer(consume(coarse_section_length), dtype=np.uint8).reshape((y_res, x_res))
-    if header[:2] == bytes([0xca, 0xac]):
+    if model == 'c201':
+        if header[-2:] != bytes([0xac,0xca]):
+            raise ValueError(f'Header end marker not found. Got header: {header[-2]:02x} {header[-1]:02x}')
+        coarse_img = np.frombuffer(consume(coarse_section_length), dtype=np.uint8).reshape((y_res, x_res))
         # 1/16th Kelvin steps
         fine_img = np.frombuffer(consume(x_res*y_res*2), dtype=np.uint16).reshape((y_res, x_res))
         fine_img = (fine_img / 16) - 273
 
         vis_jpg = Image.open(io.BytesIO(consume(jpeg_length)))
 
-    else: # 0xbaac variant
+    elif model == 'other':
+        if header[-2:] != bytes([0xab,0xba]):
+            raise ValueError(f'Header end marker not found. Got header: {header[-2]:02x} {header[-1]:02x}')
+        coarse_img = np.frombuffer(consume(coarse_section_length), dtype=np.uint8).reshape((y_res, x_res))
         # 0.1 Kelvin steps
         fine_img = np.frombuffer(consume(x_res*y_res*2), dtype=np.uint16).reshape((y_res, x_res))
         fine_img = fine_img / 10 - 273
 
         vis_jpg = Image.open(io.BytesIO(data))
+
+    else:
+        header += consume(128)
+        coarse_img = np.frombuffer(consume(coarse_section_length), dtype=np.uint8).reshape((y_res, x_res))
+        fine_img = np.frombuffer(consume(x_res*y_res*2), dtype=np.uint16).reshape((y_res, x_res))
+        fine_img = fine_img / 10 - 273
+
+        # In my example file, data now contains the JSON '{"roi":[]}' and no JPG. We ignore that.
+        vis_jpg = None
     
     return coarse_img, fine_img, vis_jpg
 
-- 
cgit