Update app.py

app.py

@@ -2,84 +2,267 @@ import gradio as gr
 import bz2
 import bson
 import struct
+from typing import List, Dict, Any
 
-PARTICLE_FORMAT = "<iii ffff f ii iiiii I"  # 14 полей = 56 байт
-FIELDS = [
-    "type", "life", "ctype",
-    "x", "y", "vx", "vy", "temp",
-    "tmp3", "tmp4", "flags", "tmp", "tmp2", "dcolour"
-]
-PARTICLE_SIZE = 128
-
-def parse_parts(binary_data):
-    count = len(binary_data) // PARTICLE_SIZE
-    particles = []
-    for i in range(count):
-        chunk = binary_data[i * PARTICLE_SIZE : i * PARTICLE_SIZE + 56]
-        try:
-            values = struct.unpack(PARTICLE_FORMAT, chunk)
-            particle = dict(zip(FIELDS, values))
-            particles.append(particle)
-        except Exception as e:
-            particles.append({"error": str(e)})
+# ---------- низкоуровневые хелперы ----------
+def _u8(buf, off) -> int:
+    return buf[off]
+def _u16(buf, off) -> int:
+    return struct.unpack_from('<H', buf, off)[0]
+def _i8(buf, off) -> int:
+    return struct.unpack_from('<b', buf, off)[0]
+
+# ---------- чтение partsPos ----------
+def read_parts_pos(partsPos: bytes, bx: int, by: int) -> List[int]:
+    """Возвращает сколько частиц в каждой клетке (длина = bx*by)"""
+    if len(partsPos) != bx * by * 3:
+        raise ValueError(f'partsPos length mismatch: expected {bx*by*3}, got {len(partsPos)}')
+    cnt_map = []
+    for i in range(bx * by):
+        base = i * 3
+        cnt = (partsPos[base] << 16) | (partsPos[base + 1] << 8) | partsPos[base + 2]
+        cnt_map.append(cnt)
+    return cnt_map
+
+# ---------- класс Particle ----------
+class Particle:
+    __slots__ = ('type', 'x', 'y', 'temp', 'life', 'tmp', 'ctype',
+                 'dcolour', 'vx', 'vy', 'tmp2', 'tmp3', 'tmp4')
+    
+    def __init__(self):
+        self.type: int = 0
+        self.x: float = 0.0
+        self.y: float = 0.0
+        self.temp: float = 0.0
+        self.life: int = 0
+        self.tmp: int = 0
+        self.ctype: int = 0
+        self.dcolour: int = 0
+        self.vx: float = 0.0
+        self.vy: float = 0.0
+        self.tmp2: int = 0
+        self.tmp3: int = 0
+        self.tmp4: int = 0
+    
+    def to_dict(self) -> Dict[str, Any]:
+        return {field: getattr(self, field) for field in self.__slots__}
+    
+    def __repr__(self):
+        return f'Particle(type={self.type} @{self.x:.1f},{self.y:.1f})'
+
+# ---------- чтение parts ----------
+def read_parts(partsData: bytes, partsPosCounts: List[int],
+               bx: int, by: int, cell_size: int = 4) -> List[Dict[str, Any]]:
+    """
+    partsData – байты из BSON['parts']
+    partsPosCounts – результат read_parts_pos
+    Возвращает список частиц как dict
+    """
+    px_per_cell = cell_size
+    particles: List[Dict[str, Any]] = []
+    off = 0
+    data_len = len(partsData)
+
+    for y in range(by):
+        for x in range(bx):
+            cnt = partsPosCounts[y * bx + x]
+            base_x, base_y = x * px_per_cell, y * px_per_cell
+            
+            for _ in range(cnt):
+                if off >= data_len - 3:
+                    raise ValueError('truncated partsData')
+                
+                # 1. type (1 байт + возможно второй)
+                ty = _u8(partsData, off)
+                off += 1
+                
+                # 2. field-descriptor 2 байта
+                fd = _u16(partsData, off)
+                off += 2
+                
+                has_fd3 = bool(fd & 0x8000)
+                if has_fd3:
+                    if off >= data_len:
+                        raise ValueError('no fd3 byte')
+                    fd3 = _u8(partsData, off)
+                    off += 1
+                    fd |= fd3 << 16
+                
+                # если type > 255 – второй байт
+                if fd & (1 << 14):
+                    ty |= _u8(partsData, off) << 8
+                    off += 1
+
+                p = Particle()
+                p.type = ty
+                p.x = float(base_x + px_per_cell // 2)
+                p.y = float(base_y + px_per_cell // 2)
+
+                # temperature
+                if fd & 1:  # 16-бит абсолютное
+                    p.temp = float(_u16(partsData, off))
+                    off += 2
+                else:  # 1 байт относительно 294.15 K
+                    delta = _i8(partsData, off)
+                    off += 1
+                    p.temp = 294.15 + delta
+
+                # life
+                if fd & (1 << 1):
+                    life = _u8(partsData, off)
+                    off += 1
+                    if fd & (1 << 2):
+                        life |= _u8(partsData, off) << 8
+                        off += 1
+                    p.life = life
+
+                # tmp
+                if fd & (1 << 3):
+                    tmp = _u8(partsData, off)
+                    off += 1
+                    if fd & (1 << 4):
+                        tmp |= _u8(partsData, off) << 8
+                        off += 1
+                        if fd & (1 << 12):
+                            tmp |= _u8(partsData, off) << 24
+                            off += 1
+                            tmp |= _u8(partsData, off) << 16
+                            off += 1
+                    p.tmp = tmp
+
+                # ctype
+                if fd & (1 << 5):
+                    ct = _u8(partsData, off)
+                    off += 1
+                    if fd & (1 << 9):
+                        ct |= _u8(partsData, off) << 24
+                        off += 1
+                        ct |= _u8(partsData, off) << 16
+                        off += 1
+                        ct |= _u8(partsData, off) << 8
+                        off += 1
+                    p.ctype = ct
+
+                # dcolour
+                if fd & (1 << 6):
+                    p.dcolour = (_u8(partsData, off) << 24) | \
+                                (_u8(partsData, off + 1) << 16) | \
+                                (_u8(partsData, off + 2) << 8) | \
+                                _u8(partsData, off + 3)
+                    off += 4
+
+                # vx, vy
+                if fd & (1 << 7):
+                    p.vx = (_u8(partsData, off) - 127.0) / 16.0
+                    off += 1
+                if fd & (1 << 8):
+                    p.vy = (_u8(partsData, off) - 127.0) / 16.0
+                    off += 1
+
+                # tmp2
+                if fd & (1 << 10):
+                    t2 = _u8(partsData, off)
+                    off += 1
+                    if fd & (1 << 11):
+                        t2 |= _u8(partsData, off) << 8
+                        off += 1
+                    p.tmp2 = t2
+
+                # tmp3/tmp4
+                if fd & (1 << 13):
+                    p.tmp3 = _u16(partsData, off)
+                    p.tmp4 = _u16(partsData, off + 2)
+                    off += 4
+                    if fd & (1 << 16):  # high halves
+                        p.tmp3 |= _u16(partsData, off) << 16
+                        p.tmp4 |= _u16(partsData, off + 2) << 16
+                        off += 4
+
+                particles.append(p.to_dict())
+    
     return particles
 
+# ---------- основная функция Gradio ----------
 def analyze_cps(file):
     if file is None:
-        return "", "", "", ""
-
+        return "", "", ""
+    
     try:
         with open(file.name, "rb") as f:
             raw_data = f.read()
-
+        
         hex_orig = ' '.join(f'{b:02x}' for b in raw_data[:256])
         hex_original = f"Оригинал (первые 256 байт):\n{hex_orig}"
-
+        
         if raw_data[:4] != b'OPS1':
-            return hex_original, "[Ошибка] Неверная сигнатура файла.", "", ""
-
+            return hex_original, "[Ошибка] Неверная сигнатура файла.", ""
+        
         compressed = raw_data[12:]
         try:
             decompressed = bz2.decompress(compressed)
         except Exception as e:
-            return hex_original, f"[Ошибка bzip2]: {e}", "", ""
-
+            return hex_original, f"[Ошибка bzip2]: {e}", ""
+        
         hex_decomp = ' '.join(f'{b:02x}' for b in decompressed[:512])
         hex_decomp_view = f"Распакованные данные (первые 512 байт):\n{hex_decomp}"
-
+        
         try:
             doc = bson.decode(decompressed)
         except Exception as e:
-            return hex_original, hex_decomp_view, f"[Ошибка BSON]: {e}", ""
-
+            return hex_original, hex_decomp_view, f"[Ошибка BSON]: {e}"
+        
         # Парсим parts
-        if 'parts' in doc:
+        if 'parts' in doc and 'partsPos' in doc and 'blockSize' in doc:
             try:
-                binary_data = bytes(doc['parts'])
-                parsed_parts = parse_parts(binary_data)
-                doc['parts'] = parsed_parts
+                # Получаем размеры блока
+                block_size = doc['blockSize']
+                if isinstance(block_size, dict):
+                    bx, by = block_size.get('x', 0), block_size.get('y', 0)
+                elif isinstance(block_size, (list, tuple)):
+                    bx, by = block_size[0], block_size[1]
+                else:
+                    bx = by = int(block_size)
+                
+                if bx <= 0 or by <= 0:
+                    raise ValueError(f"Invalid blockSize: {block_size}")
+                
+                # Читаем partsPos и parts
+                parts_pos_raw = bytes(doc['partsPos'])
+                parts_raw = bytes(doc['parts'])
+                
+                counts = read_parts_pos(parts_pos_raw, bx, by)
+                parsed_parts = read_parts(parts_raw, counts, bx, by)
+                
+                # Заменяем бинарные данные на распарсенные
+                doc['parts'] = f"[{len(parsed_parts)} particles parsed]"
+                doc['partsPos'] = f"[{len(counts)} cells]"
+                doc['_particles'] = parsed_parts  # Сохраняем отдельно
+                
             except Exception as e:
-                doc['parts'] = f"[Ошибка парсинга parts]: {e}"
-
+                doc['parts'] = f"[Ошибка парсинга]: {e}"
+                doc['partsPos'] = f"[Ошибка парсинга]: {e}"
+        
         return hex_original, hex_decomp_view, str(doc)
-
+    
     except Exception as e:
-        return f"[Ошибка]: {e}", "", "", ""
+        return f"[Ошибка]: {e}", "", ""
 
+# ---------- Gradio интерфейс ----------
 with gr.Blocks() as demo:
-    gr.Markdown("## Распаковка .cps с красивым parts")
+    gr.Markdown("## Распаковка .cps (The Powder Toy) с полным парсингом частиц")
+    
     file_input = gr.File(label="Загрузите .cps файл", file_types=[".cps"])
-
+    
     hex_orig_out = gr.Textbox(label="HEX оригинала", lines=5)
     hex_decomp_out = gr.Textbox(label="HEX распакованных данных", lines=10)
-    bson_parsed = gr.Textbox(label="Распарсенный BSON с parts", lines=30)
-
+    bson_parsed = gr.Textbox(label="Распарсенный BSON", lines=30)
+    
     file_input.change(
         fn=analyze_cps,
         inputs=file_input,
         outputs=[hex_orig_out, hex_decomp_out, bson_parsed]
     )
-
+    
     file_input.clear(
         fn=lambda: ("", "", ""),
         outputs=[hex_orig_out, hex_decomp_out, bson_parsed]