添加缩略图提取逻辑

heic_file.py

@@ -95,7 +95,11 @@ class HEICFile:
         primary_id = self.get_primary_item_id()
         if not primary_id: return None
         if not self._iref_box: return None
-        return self._iref_box.references.get(primary_id)
+        
+        # 'dimg' (derived image) 是网格布局的标准引用类型
+        if 'dimg' in self._iref_box.references:
+            return self._iref_box.references['dimg'].get(primary_id)
+        return None
 
     def get_image_size(self, item_id: int) -> tuple[int, int] | None:
         if not (self._iprp_box and self._iprp_box.ipma and self._iprp_box.ipco): return None
@@ -152,4 +156,48 @@ class HEICFile:
             with open(self.filepath, 'rb') as f:
                 f.seek(offset)
                 return f.read()
-        return None
+        return None
+
+    def get_thumbnail_data(self) -> bytes | None:
+        """
+        查找并提取缩略图图像数据 (如果存在)。
+        它会查找一个 'iref' 盒子，其中包含 'thmb' (thumbnail) 引用。
+        """
+        print("Attempting to extract thumbnail data...")
+        primary_id = self.get_primary_item_id()
+        if not primary_id:
+            print("Error: Could not determine primary item ID.")
+            return None
+        
+        if not self._iref_box:
+            print("Info: No 'iref' box found, cannot search for thumbnail.")
+            return None
+
+        # 检查 'thmb' (thumbnail) 引用类型是否存在
+        if 'thmb' not in self._iref_box.references:
+            print("Info: No 'thmb' references found in 'iref' box.")
+            return None
+        
+        # 检查主图像是否有与之关联的缩略图
+        if primary_id not in self._iref_box.references['thmb']:
+            print(f"Info: Primary item ID {primary_id} has no 'thmb' reference.")
+            return None
+
+        thumbnail_ids = self._iref_box.references['thmb'][primary_id]
+        if not thumbnail_ids:
+            print(f"Info: Primary item ID {primary_id} has 'thmb' reference, but no target IDs.")
+            return None
+
+        # 获取第一个缩略图 ID
+        thumbnail_id = thumbnail_ids[0]
+        print(f"Found thumbnail reference: Primary ID {primary_id} -> Thumbnail ID {thumbnail_id}")
+        
+        # 使用现有的 get_item_data 来提取它
+        thumbnail_data = self.get_item_data(thumbnail_id)
+        
+        if thumbnail_data:
+            print(f"Successfully extracted thumbnail data (Item ID {thumbnail_id}).")
+            return thumbnail_data
+        else:
+            print(f"Error: Failed to get data for thumbnail item ID {thumbnail_id}.")
+            return None

heic_types.py

@@ -206,8 +206,10 @@ class ItemReferenceEntry:
 class ItemReferenceBox(Box):
     def __init__(self, size: int, box_type: str, offset: int, raw_data: bytes):
         super().__init__(size, box_type, offset, raw_data)
-        self.references: dict[int, list[int]] = {}
+        # 结构: {'thmb': {from_id: [to_id, ...]}, 'dimg': {from_id: [to_id, ...]}}
+        self.references: dict[str, dict[int, list[int]]] = {}
         self._parse_references()
+        
     def _parse_references(self):
         stream = self.raw_data
         version_flags = struct.unpack('>I', stream[:4])[0]
@@ -215,17 +217,29 @@ class ItemReferenceBox(Box):
         pos = 4
         item_id_size = 4 if version == 1 else 2
         while pos < len(stream):
-            if pos + 8 > len(stream): break
+            if pos + 8 > len(stream): break # 需要空间来读取 size 和 type
             ref_box_size = struct.unpack('>I', stream[pos:pos+4])[0]
+            ref_box_type = stream[pos+4:pos+8].decode('ascii', errors='ignore')
+            
             if pos + ref_box_size > len(stream): break
+            
+            # 为这种引用类型创建字典 (如果它还不存在)
+            if ref_box_type not in self.references:
+                self.references[ref_box_type] = {}
+            
             if item_id_size == 4:
+                if pos + 12 > len(stream): break
                 from_item_id = struct.unpack('>I', stream[pos+8:pos+12])[0]
                 ref_count_pos = pos + 12
             else:
+                if pos + 10 > len(stream): break
                 from_item_id = struct.unpack('>H', stream[pos+8:pos+10])[0]
                 ref_count_pos = pos + 10
+                
+            if ref_count_pos + 2 > len(stream): break
             reference_count = struct.unpack('>H', stream[ref_count_pos:ref_count_pos+2])[0]
             to_ids_pos = ref_count_pos + 2
+            
             to_item_ids = []
             for _ in range(reference_count):
                 if to_ids_pos + item_id_size > len(stream): break
@@ -235,5 +249,6 @@ class ItemReferenceBox(Box):
                     to_id = struct.unpack('>H', stream[to_ids_pos:to_ids_pos+2])[0]
                 to_item_ids.append(to_id)
                 to_ids_pos += item_id_size
-            self.references[from_item_id] = to_item_ids
+                
+            self.references[ref_box_type][from_item_id] = to_item_ids
             pos += ref_box_size

main.py

@@ -40,6 +40,23 @@ def analyze_and_reconstruct(filename: str):
     else:
         print("Failed to reconstruct image.")
 
+    # 3. 尝试提取并保存缩略图
+    thumbnail_data = heic_file.get_thumbnail_data()
+    if thumbnail_data:
+        try:
+            # 缩略图数据通常是 JPEG 或 HEIC 格式
+            # 我们先假设它是 JPEG 并以此后缀保存
+            thumb_filename = f"{base_filename}_thumbnail.jpg"
+            with open(thumb_filename, "wb") as f_thumb:
+                f_thumb.write(thumbnail_data)
+            print(f"Successfully saved thumbnail data to: {thumb_filename}")
+            print(f" (Note: This file might be a JPEG or a small HEIC. Try opening it.)")
+        except Exception as e:
+            print(f"Could not save thumbnail file: {e}")
+    else:
+        print("No thumbnail data found or extracted.")
+
+
     print("\n" + "="*40 + "\n")
 
 def main():