添加pyheif

heic_file.py

@@ -1,26 +1,42 @@
 # pyheic_struct/heic_file.py
 
 import os
+import math
+# 1. 导入新的库
+import pillow_heif
+from io import BytesIO
+from dataclasses import dataclass
+from PIL import Image
+
 from base import Box
 from parser import parse_boxes
 from heic_types import (
     ItemLocationBox, PrimaryItemBox, ItemInfoBox, ItemPropertiesBox,
-    ImageSpatialExtentsBox, ItemReferenceBox # <-- Import the new class
+    ImageSpatialExtentsBox, ItemReferenceBox
 )
-
 from handlers.base_handler import VendorHandler
 from handlers.apple_handler import AppleHandler
 from handlers.samsung_handler import SamsungHandler
 
+@dataclass
+class Grid:
+    rows: int
+    columns: int
+    output_width: int
+    output_height: int
+
 class HEICFile:
     def __init__(self, filepath: str):
         self.filepath = filepath
+        # 注册 HEIF/HEIC 文件格式解码器
+        pillow_heif.register_heif_opener()
+        
         self._iloc_box: ItemLocationBox | None = None
         self._ftyp_box: Box | None = None
         self._iinf_box: ItemInfoBox | None = None
         self._pitm_box: PrimaryItemBox | None = None
         self._iprp_box: ItemPropertiesBox | None = None
-        self._iref_box: ItemReferenceBox | None = None # <-- Add new instance variable
+        self._iref_box: ItemReferenceBox | None = None
         self.handler: VendorHandler | None = None
 
         with open(self.filepath, 'rb') as f:
@@ -30,97 +46,85 @@ class HEICFile:
             self._detect_vendor()
 
     def _find_essential_boxes(self, boxes: list[Box]):
-        """Recursively search for essential boxes."""
         for box in boxes:
-            if isinstance(box, ItemLocationBox):
-                self._iloc_box = box
-            if isinstance(box, ItemInfoBox):
-                self._iinf_box = box
-            if isinstance(box, PrimaryItemBox):
-                self._pitm_box = box
-            if isinstance(box, ItemPropertiesBox):
-                self._iprp_box = box
-            if isinstance(box, ItemReferenceBox): # <-- Find and store the 'iref' box
-                self._iref_box = box
-            if box.type == 'ftyp':
-                self._ftyp_box = box
+            if isinstance(box, ItemLocationBox): self._iloc_box = box
+            if isinstance(box, ItemInfoBox): self._iinf_box = box
+            if isinstance(box, PrimaryItemBox): self._pitm_box = box
+            if isinstance(box, ItemPropertiesBox): self._iprp_box = box
+            if isinstance(box, ItemReferenceBox): self._iref_box = box
+            if box.type == 'ftyp': self._ftyp_box = box
+            if box.children: self._find_essential_boxes(box.children)
 
-            if box.children:
-                self._find_essential_boxes(box.children)
+    # --- 2. 使用 pillow-heif 大大简化的图像重建方法 ---
+    def reconstruct_primary_image(self) -> Image.Image | None:
+        """
+        使用 pillow-heif 重建主图像，它能自动处理网格图像。
+        """
+        try:
+            print("Reconstructing primary image using pillow-heif...")
+            # pillow-heif 让我们能像打开普通图片一样直接打开 HEIC 文件
+            image = Image.open(self.filepath)
+            
+            # 确保图像数据被加载
+            image.load() 
+            
+            print("Successfully reconstructed image.")
+            return image
+        except Exception as e:
+            print(f"Failed to reconstruct image with pillow-heif: {e}")
+            return None
 
-    # --- NEW Method for Task 3.2 ---
-    def get_grid_layout(self) -> list[int] | None:
-        """
-        Uses 'iref' and 'pitm' to find the tile IDs for the primary grid image.
-        Returns a list of item IDs that make up the grid.
-        """
+    def get_primary_image_grid(self) -> Grid | None:
         primary_id = self.get_primary_item_id()
-        if not primary_id:
-            return None
+        if not primary_id: return None
+        full_size = self.get_image_size(primary_id)
+        if not full_size: return None
+        full_width, full_height = full_size
+        grid_tiles = self.get_grid_layout()
+        if not grid_tiles: return None
+        first_tile_id = grid_tiles[0]
+        tile_size = self.get_image_size(first_tile_id)
+        if not tile_size: return None
+        tile_width, tile_height = tile_size
+        if tile_width == 0 or tile_height == 0: return None
+        columns = math.ceil(full_width / tile_width)
+        rows = math.ceil(full_height / tile_height)
+        return Grid(rows=rows, columns=columns, output_width=full_width, output_height=full_height)
 
-        if not self._iref_box:
-            print("Warning: 'iref' box not found. Cannot determine grid layout.")
-            return None
-        
-        # Find the reference entry that originates from the primary item ID
-        grid_tile_ids = self._iref_box.references.get(primary_id)
-
-        if not grid_tile_ids:
-            print(f"Warning: No grid reference found for primary item ID {primary_id}.")
-            return None
-
-        return grid_tile_ids
+    def get_grid_layout(self) -> list[int] | None:
+        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)
 
     def get_image_size(self, item_id: int) -> tuple[int, int] | None:
-        """
-        Gets the width and height of a given image item ID.
-        """
-        if not self._iprp_box or not self._iprp_box.ipma or not self._iprp_box.ipco:
-            print("Warning: 'iprp' box or its sub-boxes ('ipma', 'ipco') not found.")
-            return None
-        
-        if item_id not in self._iprp_box.ipma.entries:
-            print(f"Warning: No property association found for item ID {item_id}.")
-            return None
-        
+        if not (self._iprp_box and self._iprp_box.ipma and self._iprp_box.ipco): return None
+        if item_id not in self._iprp_box.ipma.entries: return None
         item_associations = self._iprp_box.ipma.entries[item_id].associations
-
         for assoc in item_associations:
-            property_index = assoc - 1 # Property indices are 1-based
+            property_index = assoc - 1
             if 0 <= property_index < len(self._iprp_box.ipco.children):
                 prop = self._iprp_box.ipco.children[property_index]
                 if isinstance(prop, ImageSpatialExtentsBox):
                     return (prop.image_width, prop.image_height)
-        
-        print(f"Warning: 'ispe' property not found for item ID {item_id}.")
         return None
 
     def get_primary_item_id(self) -> int | None:
-        if not self._pitm_box:
-            print("Warning: 'pitm' box not found. Cannot determine primary item.")
-            return None
+        if not self._pitm_box: return None
         return self._pitm_box.item_id
 
     def list_items(self):
-        if not self._iinf_box:
-            print("Warning: 'iinf' box not found. Cannot list items.")
-            return
-        
+        if not self._iinf_box: return
         print("Available items in HEIC file:")
-        for entry in self._iinf_box.entries:
-            print(f"  - {entry}")
+        for entry in self._iinf_box.entries: print(f"  - {entry}")
 
     def _detect_vendor(self):
         if self._ftyp_box:
             compatible_brands = self._ftyp_box.raw_data[4:].decode('ascii', errors='ignore')
-            if 'samsung' in compatible_brands.lower() or 'mpvd' in [b.type for b in self.boxes]:
-                 self.handler = SamsungHandler()
-                 return
-            if 'apple' in compatible_brands.lower() or 'MiHB' in compatible_brands:
-                 self.handler = AppleHandler()
-                 return
-        
-        self.handler = VendorHandler()
+            if 'samsung' in compatible_brands.lower(): self.handler = SamsungHandler()
+            elif 'apple' in compatible_brands.lower() or 'MiHB' in compatible_brands: self.handler = AppleHandler()
+            else: self.handler = VendorHandler()
+        else: self.handler = VendorHandler()
 
     def get_item_data(self, item_id: int) -> bytes | None:
         if not self._iloc_box:
@@ -130,19 +134,22 @@ class HEICFile:
         if not location:
             print(f"Error: Item with ID {item_id} not found in 'iloc' box.")
             return None
+        if not location.extents:
+             print(f"Warning: Item with ID {item_id} has no extents.")
+             return b''
+
+        data_chunks = []
         with open(self.filepath, 'rb') as f:
-            f.seek(location.offset)
-            data = f.read(location.length)
-            return data
+            for offset, length in location.extents:
+                f.seek(offset)
+                data_chunks.append(f.read(length))
+        return b''.join(data_chunks)
 
     def get_motion_photo_data(self) -> bytes | None:
-        if not self.handler:
-            self._detect_vendor()
+        if not self.handler: self._detect_vendor()
         offset = self.handler.find_motion_photo_offset(self)
         if offset is not None:
             with open(self.filepath, 'rb') as f:
                 f.seek(offset)
-                # The actual video data could be the rest of the file or a specific length
-                # For now, we assume it's a reasonable chunk.
                 return f.read()
         return None

heic_types.py

@@ -3,50 +3,100 @@
 import struct
 from base import Box
 
-# --- Existing classes (unchanged) ---
-
+# --- A new helper class to handle multiple data locations ('extents') ---
 class ItemLocation:
-    def __init__(self, item_id, offset, length):
+    def __init__(self, item_id):
         self.item_id = item_id
-        self.offset = offset
-        self.length = length
-    def __repr__(self):
-        return f"<ItemLocation ID={self.item_id} offset={self.offset} length={self.length}>"
+        self.extents = [] # Will be a list of (offset, length) tuples
 
+    def __repr__(self):
+        total_length = sum(ext[1] for ext in self.extents)
+        return f"<ItemLocation ID={self.item_id} extents={len(self.extents)} total_size={total_length}>"
+
+# --- REWRITTEN ItemLocationBox to correctly parse all item locations ---
 class ItemLocationBox(Box):
     def __init__(self, size: int, box_type: str, offset: int, raw_data: bytes):
         super().__init__(size, box_type, offset, raw_data)
         self.locations = []
         self._parse_locations()
+        
     def _parse_locations(self):
         stream = self.raw_data
+        
         version_flags = struct.unpack('>I', stream[:4])[0]
+        version = version_flags >> 24
+        
         sizes = struct.unpack('>H', stream[4:6])[0]
         offset_size = (sizes >> 12) & 0x0F
         length_size = (sizes >> 8) & 0x0F
         base_offset_size = (sizes >> 4) & 0x0F
-        item_count = struct.unpack('>H', stream[6:8])[0]
-        current_pos = 8
+        
+        index_size = 0
+        if version == 1 or version == 2:
+            index_size = sizes & 0x0F
+        
+        item_count = 0
+        if version < 2:
+            item_count = struct.unpack('>H', stream[6:8])[0]
+            current_pos = 8
+        else: # version == 2
+            item_count = struct.unpack('>I', stream[6:10])[0]
+            current_pos = 10
+
         for _ in range(item_count):
-            item_id = struct.unpack('>H', stream[current_pos : current_pos+2])[0]
-            current_pos += 2
-            current_pos += 2 # Skip 2 bytes of construction_method
-            current_pos += base_offset_size # Skip base_offset
+            item_id = 0
+            # Item ID size depends on the box version
+            if version < 2:
+                if current_pos + 2 > len(stream): break
+                item_id = struct.unpack('>H', stream[current_pos : current_pos+2])[0]
+                current_pos += 2
+            else: # version == 2
+                if current_pos + 4 > len(stream): break
+                item_id = struct.unpack('>I', stream[current_pos : current_pos+4])[0]
+                current_pos += 4
+
+            if (version == 1 or version == 2) and current_pos + 2 <= len(stream):
+                current_pos += 2 # Skip construction method
+
+            if current_pos + 2 > len(stream): break
+            current_pos += 2 # Skip data_reference_index
+            
+            base_offset = 0
+            if base_offset_size > 0:
+                if current_pos + base_offset_size > len(stream): break
+                base_offset = self._read_int(stream, current_pos, base_offset_size)
+                current_pos += base_offset_size
+
+            if current_pos + 2 > len(stream): break
             extent_count = struct.unpack('>H', stream[current_pos : current_pos+2])[0]
             current_pos += 2
-            if extent_count > 0:
-                offset = self._read_int(stream, current_pos, offset_size)
+            
+            loc = ItemLocation(item_id)
+            for __ in range(extent_count):
+                if (version == 1 or version == 2) and index_size > 0:
+                     if current_pos + index_size > len(stream): break
+                     current_pos += index_size # Skip extent_index
+
+                extent_offset = self._read_int(stream, current_pos, offset_size)
                 current_pos += offset_size
-                length = self._read_int(stream, current_pos, length_size)
+
+                extent_length = self._read_int(stream, current_pos, length_size)
                 current_pos += length_size
-                self.locations.append(ItemLocation(item_id, offset, length))
+                
+                loc.extents.append((base_offset + extent_offset, extent_length))
+            self.locations.append(loc)
+
     def _read_int(self, data, pos, size):
+        if pos + size > len(data): return 0
+        if size == 0: return 0
         if size == 1: return data[pos]
         if size == 2: return struct.unpack('>H', data[pos:pos+2])[0]
         if size == 4: return struct.unpack('>I', data[pos:pos+4])[0]
         if size == 8: return struct.unpack('>Q', data[pos:pos+8])[0]
         return 0
 
+# --- All other classes below remain unchanged ---
+
 class ItemInfoEntry:
     def __init__(self, item_id, item_type, item_name):
         self.item_id = item_id
@@ -61,12 +111,10 @@ class ItemInfoBox(Box):
         self.entries: list[ItemInfoEntry] = []
         self._parse_entries()
     def _parse_entries(self):
-        # Assumes version 0 of 'iinf' box
         if len(self.raw_data) < 6: return
         item_count = struct.unpack('>H', self.raw_data[4:6])[0]
         for infe_box in self.children:
             if infe_box.type == 'infe':
-                # Assumes version 2 of 'infe' box
                 if len(infe_box.raw_data) < 12: continue
                 item_id = struct.unpack('>H', infe_box.raw_data[4:6])[0]
                 item_type = infe_box.raw_data[8:12].decode('ascii').strip('\x00')
@@ -80,12 +128,10 @@ class PrimaryItemBox(Box):
         self.item_id: int = 0
         self._parse_item_id()
     def _parse_item_id(self):
-        # Assumes version 0 of 'pitm' box
         if len(self.raw_data) < 6: return
         self.item_id = struct.unpack('>H', self.raw_data[4:6])[0]
 
 class ImageSpatialExtentsBox(Box):
-    """'ispe' box, contains image dimensions."""
     def __init__(self, size: int, box_type: str, offset: int, raw_data: bytes):
         super().__init__(size, box_type, offset, raw_data)
         self.image_width = struct.unpack('>I', self.raw_data[4:8])[0]
@@ -94,7 +140,6 @@ class ImageSpatialExtentsBox(Box):
         return f"<ImageSpatialExtentsBox width={self.image_width} height={self.image_height}>"
 
 class ItemPropertyAssociationEntry:
-    """Helper class for an entry in the 'ipma' box."""
     def __init__(self, item_id, association_count):
         self.item_id = item_id
         self.association_count = association_count
@@ -103,7 +148,6 @@ class ItemPropertyAssociationEntry:
         return f"<ItemPropertyAssociationEntry item_id={self.item_id} associations={self.associations}>"
 
 class ItemPropertyAssociationBox(Box):
-    """'ipma' box, maps items to their properties."""
     def __init__(self, size: int, box_type: str, offset: int, raw_data: bytes):
         super().__init__(size, box_type, offset, raw_data)
         self.entries: dict[int, ItemPropertyAssociationEntry] = {}
@@ -138,31 +182,21 @@ class ItemPropertyAssociationBox(Box):
             self.entries[item_id] = entry
 
 class ItemPropertyContainerBox(Box):
-    """'ipco' box. Just a container for property boxes like 'ispe'."""
     pass
 
 class ItemPropertiesBox(Box):
-    """'iprp' box. Container for 'ipco' and 'ipma'."""
     @property
     def ipco(self) -> ItemPropertyContainerBox | None:
-        """Finds and returns the 'ipco' child box if it exists."""
         for child in self.children:
-            if isinstance(child, ItemPropertyContainerBox):
-                return child
+            if isinstance(child, ItemPropertyContainerBox): return child
         return None
-
     @property
     def ipma(self) -> ItemPropertyAssociationBox | None:
-        """Finds and returns the 'ipma' child box if it exists."""
         for child in self.children:
-            if isinstance(child, ItemPropertyAssociationBox):
-                return child
+            if isinstance(child, ItemPropertyAssociationBox): return child
         return None
 
-# --- NEW classes for Task 3.1 ('iref') ---
-
 class ItemReferenceEntry:
-    """Helper class for a single reference in the 'iref' box."""
     def __init__(self, from_id, to_ids):
         self.from_item_id = from_id
         self.to_item_ids = to_ids
@@ -170,54 +204,36 @@ class ItemReferenceEntry:
         return f"<ItemReferenceEntry from={self.from_item_id} to={self.to_item_ids}>"
 
 class ItemReferenceBox(Box):
-    """'iref' box, describes relationships between items."""
     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]] = {}
         self._parse_references()
-
     def _parse_references(self):
         stream = self.raw_data
         version_flags = struct.unpack('>I', stream[:4])[0]
         version = version_flags >> 24
         pos = 4
-        
         item_id_size = 4 if version == 1 else 2
-
-        # The 'iref' box contains one or more SingleItemTypeReferenceBox(es)
-        # We parse them in a loop until we run out of data.
         while pos < len(stream):
-            # Each sub-box has its own size and type
             if pos + 8 > len(stream): break
             ref_box_size = struct.unpack('>I', stream[pos:pos+4])[0]
-            ref_box_type = stream[pos+4:pos+8].decode('ascii')
-            
             if pos + ref_box_size > len(stream): break
-
-            # We are interested in 'grid', but could parse others like 'thmb'
-            # The structure is the same for SingleItemTypeReferenceBox
-            
             if item_id_size == 4:
                 from_item_id = struct.unpack('>I', stream[pos+8:pos+12])[0]
                 ref_count_pos = pos + 12
-            else: # item_id_size == 2
+            else:
                 from_item_id = struct.unpack('>H', stream[pos+8:pos+10])[0]
                 ref_count_pos = pos + 10
-                
             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
                 if item_id_size == 4:
                     to_id = struct.unpack('>I', stream[to_ids_pos:to_ids_pos+4])[0]
-                else: # item_id_size == 2
+                else:
                     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
-            
-            # Move to the next reference box within 'iref'
             pos += ref_box_size

main.py

@@ -1,39 +1,54 @@
 # pyheic_struct/main.py
 
 from heic_file import HEICFile
+import os
 
-def analyze_file(filename: str):
+def analyze_and_reconstruct(filename: str):
     print(f"--- Analyzing {filename} ---")
+    
+    # 检查文件是否存在
+    if not os.path.exists(filename):
+        print(f"Error: File not found at {filename}")
+        print("="*40 + "\n")
+        return
+
     heic_file = HEICFile(filename)
 
-    # List all available items in the file
-    heic_file.list_items()
-
-    # Get the ID of the primary image
-    primary_id = heic_file.get_primary_item_id()
-    if primary_id:
-        print(f"The primary item ID is: {primary_id}")
-
-        # Get the primary image dimensions
-        size = heic_file.get_image_size(primary_id)
-        if size:
-            print(f"Primary image dimensions (WxH): {size[0]} x {size[1]}")
+    # 重建主图像
+    reconstructed_image = heic_file.reconstruct_primary_image()
+    
+    if reconstructed_image:
+        # 获取不带扩展名的基本文件名
+        base_filename = os.path.splitext(filename)[0]
         
-        # --- Use our new feature! ---
-        grid_layout = heic_file.get_grid_layout()
-        if grid_layout:
-            print(f"Primary image is a grid composed of these item IDs: {grid_layout}")
-            # Optional: Get size of the first tile to see the difference
-            first_tile_id = grid_layout[0]
-            tile_size = heic_file.get_image_size(first_tile_id)
-            if tile_size:
-                print(f"  - Size of the first tile (ID {first_tile_id}): {tile_size[0]} x {tile_size[1]}")
+        # 1. 保存为高质量的有损格式 (JPEG)
+        try:
+            jpeg_filename = f"{base_filename}_reconstructed.jpg"
+            reconstructed_image.save(jpeg_filename, "JPEG", quality=95)
+            print(f"Successfully saved lossy JPEG to: {jpeg_filename}")
+        except Exception as e:
+            print(f"Could not save JPEG file: {e}")
+
+        # 2. 保存为无损格式 (PNG)
+        try:
+            png_filename = f"{base_filename}_reconstructed.png"
+            reconstructed_image.save(png_filename, "PNG")
+            print(f"Successfully saved lossless PNG to: {png_filename}")
+        except Exception as e:
+            print(f"Could not save PNG file: {e}")
+            
+    else:
+        print("Failed to reconstruct image.")
 
     print("\n" + "="*40 + "\n")
 
 def main():
-    analyze_file('apple.HEIC')
-    analyze_file('samsung.heic')
+    # 确保文件路径正确
+    apple_file = 'apple.HEIC'
+    samsung_file = 'samsung.heic'
+    
+    analyze_and_reconstruct(apple_file)
+    analyze_and_reconstruct(samsung_file)
 
 if __name__ == "__main__":
     main()