diff --git a/.DS_Store b/.DS_Store index b015eee..1799903 100644 Binary files a/.DS_Store and b/.DS_Store differ diff --git a/__pycache__/heic_file.cpython-314.pyc b/__pycache__/heic_file.cpython-314.pyc index 0b51bbc..1650258 100644 Binary files a/__pycache__/heic_file.cpython-314.pyc and b/__pycache__/heic_file.cpython-314.pyc differ diff --git a/__pycache__/heic_types.cpython-314.pyc b/__pycache__/heic_types.cpython-314.pyc index 524e191..50a5690 100644 Binary files a/__pycache__/heic_types.cpython-314.pyc and b/__pycache__/heic_types.cpython-314.pyc differ diff --git a/apple_reconstructed.jpg b/apple_reconstructed.jpg new file mode 100644 index 0000000..32a0062 Binary files /dev/null and b/apple_reconstructed.jpg differ diff --git a/apple_reconstructed.png b/apple_reconstructed.png new file mode 100644 index 0000000..e77ecc5 Binary files /dev/null and b/apple_reconstructed.png differ diff --git a/heic_file.py b/heic_file.py index cc7b71b..203ea85 100644 --- a/heic_file.py +++ b/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 \ No newline at end of file diff --git a/heic_types.py b/heic_types.py index a7ac1bd..b1b4fc0 100644 --- a/heic_types.py +++ b/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"" + 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"" + +# --- 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"" 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"" 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"" 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 \ No newline at end of file diff --git a/main.py b/main.py index 0c4c124..658f173 100644 --- a/main.py +++ b/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() \ No newline at end of file diff --git a/samsung_reconstructed.jpg b/samsung_reconstructed.jpg new file mode 100644 index 0000000..4a72b74 Binary files /dev/null and b/samsung_reconstructed.jpg differ diff --git a/samsung_reconstructed.png b/samsung_reconstructed.png new file mode 100644 index 0000000..c95da60 Binary files /dev/null and b/samsung_reconstructed.png differ