Update app.py

app.py

@@ -52,50 +52,39 @@ def to_jsonable(obj):
         },
     )
     
-def clean_ppstructure_result(result):
-    """
-    Recursively converts numpy arrays and custom objects in the PPStructure result
-    to standard Python lists and dictionaries that are JSON-safe.
-    """
+def clean_ppstructure_result(result, exclude_keys=("img",), include_img=False, img_mode="shape"):
     if isinstance(result, list):
-        # Case 1: List - Recurse through all elements
-        return [clean_ppstructure_result(item) for item in result]
-
-    # Case 2: Dictionary or Custom Object
-    # Custom objects (with a __dict__) need to be converted to dicts before processing.
-    if isinstance(result, dict) or hasattr(result, '__dict__'):
-        if hasattr(result, '__dict__'):
-            # It's a custom class instance (e.g., from PPStructure), get its attributes as a dict
-            items = vars(result).items()
-        else:
-            # It's a standard dictionary
-            items = result.items()
-
-        cleaned_dict = {}
+        return [clean_ppstructure_result(item, exclude_keys, include_img, img_mode) for item in result]
+
+    if isinstance(result, dict) or hasattr(result, "__dict__"):
+        items = vars(result).items() if hasattr(result, "__dict__") else result.items()
+        cleaned = {}
         for key, value in items:
-            # Exclude raw image data, as it's too large and unnecessary for JSON
-            if key == 'img':
-                 continue
-            
-            cleaned_dict[key] = clean_ppstructure_result(value)
-        return cleaned_dict
-    
-    # Case 3: Numpy array - Convert to standard Python list
-    elif isinstance(result, np.ndarray):
+            if key in exclude_keys:
+                if include_img and key == "img":
+                    if img_mode == "shape" and isinstance(value, np.ndarray):
+                        cleaned["img_shape"] = value.shape
+                    elif img_mode == "base64" and isinstance(value, np.ndarray):
+                        ok, buf = cv2.imencode(".png", value)
+                        if ok:
+                            cleaned["img_base64"] = base64.b64encode(buf).decode("ascii")
+                continue
+            cleaned[key] = clean_ppstructure_result(value, exclude_keys, include_img, img_mode)
+        return cleaned
+
+    if isinstance(result, tuple) or isinstance(result, set):
+        return [clean_ppstructure_result(v, exclude_keys, include_img, img_mode) for v in result]
+
+    if isinstance(result, np.ndarray):
         return result.tolist()
-        
-    # Case 4: Numpy scalar types - Convert to standard Python primitives
-    # These base classes cover all int and float variants (int16, float32, etc.)
-    elif isinstance(result, np.floating):
+    if isinstance(result, np.floating):
         return float(result)
-    elif isinstance(result, np.integer):
+    if isinstance(result, np.integer):
         return int(result)
-    elif isinstance(result, np.bool_):
+    if isinstance(result, np.bool_):
         return bool(result)
-        
-    # Case 5: Standard Python primitive or unhandled type - return as is
-    else:
-        return result
+
+    return result
 
 # --- API Endpoints ---
 
@@ -143,8 +132,8 @@ def run_structure_analysis(ocr_input: OcrInput):
 
         print("** Structure Analysis Complete. Converting to JSON-safe **")
         json_safe_result = clean_ppstructure_result(raw_result)
-        #return {"result": json_safe_result}
-        return {"result": to_jsonable(raw_result)}
+        return {"result": json_safe_result}
+        #return {"result": to_jsonable(raw_result)}
 
     except HTTPException:
         raise