Format response

app.py

@@ -3,6 +3,7 @@ import json
 import argparse
 import operator
 import gradio as gr
+import torch
 import torchvision
 from typing import Tuple, Dict
 from facetorch import FaceAnalyzer
@@ -23,6 +24,27 @@ args = parser.parse_args()
 cfg = OmegaConf.load(args.path_conf)
 analyzer = FaceAnalyzer(cfg.analyzer)
 
+
+def tensor_to_list(tensor: torch.Tensor) -> list:
+    return tensor.tolist()
+
+def dataclass_to_dict(obj):
+    if hasattr(obj, "__dataclass_fields__"):
+        return asdict(obj)
+    return obj
+
+def image_data_to_json(image_data: ImageData) -> str:
+    # Convert tensors to lists
+    image_data.img = tensor_to_list(image_data.img)
+    image_data.tensor = tensor_to_list(image_data.tensor)
+    
+    # Convert dataclass to dictionary
+    data_dict = dataclass_to_dict(image_data)
+    
+    # Convert dictionary to JSON string
+    json_str = json.dumps(data_dict, indent=4, default=dataclass_to_dict)
+    return json_str
+
 def gen_sim_dict_str(response: ImageData, pred_name: str = "verify", index: int = 0)-> str:     
     if len(response.faces) > 0:
         base_emb = response.faces[index].preds[pred_name].logits
@@ -34,7 +56,6 @@ def gen_sim_dict_str(response: ImageData, pred_name: str = "verify", index: int
         
     return sim_dict_sort_str
 
-
 def inference(path_image: str) -> Tuple:
     response = analyzer.run(
         path_image=path_image,
@@ -50,7 +71,7 @@ def inference(path_image: str) -> Tuple:
     fer_dict_str = str({face.indx: face.preds["fer"].label for face in response.faces})
     au_dict_str = str({face.indx: face.preds["au"].other["multi"] for face in response.faces})
     deepfake_dict_str = str({face.indx: face.preds["deepfake"].label for face in response.faces})
-    response_str = str(response)
+    response_str = image_data_to_json(response)
     
     sim_dict_str_embed = gen_sim_dict_str(response, pred_name="embed", index=0)
     sim_dict_str_verify = gen_sim_dict_str(response, pred_name="verify", index=0)