include label_info in API results

microsoft · Dec 27, 2024 · aed239a · aed239a
1 parent e6f4078
commit aed239a
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/pe/api/image/improved_diffusion_api.py b/pe/api/image/improved_diffusion_api.py
@@ -153,10 +153,12 @@ def random_api(self, label_info, num_samples):
             {
                 IMAGE_DATA_COLUMN_NAME: list(samples),
                 IMAGE_MODEL_LABEL_COLUMN_NAME: list(labels),
+                LABEL_ID_COLUMN_NAME: 0,
             }
         )
+        metadata = {"label_info": [label_info]}
         execution_logger.info(f"RANDOM API: finished creating {num_samples} samples for label {label_name}")
-        return Data(data_frame=data_frame)
+        return Data(data_frame=data_frame, metadata=metadata)
 
     def variation_api(self, syn_data):
         """Generating variations of the synthetic data.
@@ -200,6 +202,7 @@ def variation_api(self, syn_data):
             {
                 IMAGE_DATA_COLUMN_NAME: list(variations),
                 IMAGE_MODEL_LABEL_COLUMN_NAME: list(labels),
+                LABEL_ID_COLUMN_NAME: syn_data.data_frame[LABEL_ID_COLUMN_NAME].values,
             }
         )
         if LABEL_ID_COLUMN_NAME in syn_data.data_frame.columns:

diff --git a/pe/api/image/stable_diffusion_api.py b/pe/api/image/stable_diffusion_api.py
@@ -152,10 +152,12 @@ def random_api(self, label_info, num_samples):
             {
                 IMAGE_DATA_COLUMN_NAME: list(images),
                 IMAGE_PROMPT_COLUMN_NAME: prompt,
+                LABEL_ID_COLUMN_NAME: 0,
             }
         )
+        metadata = {"label_info": [label_info]}
         execution_logger.info(f"RANDOM API: finished creating {num_samples} samples for label {label_name}")
-        return Data(data_frame=data_frame)
+        return Data(data_frame=data_frame, metadata=metadata)
 
     def variation_api(self, syn_data):
         """Generating variations of the synthetic data.
@@ -199,6 +201,7 @@ def variation_api(self, syn_data):
             {
                 IMAGE_DATA_COLUMN_NAME: list(variations),
                 IMAGE_PROMPT_COLUMN_NAME: prompts,
+                LABEL_ID_COLUMN_NAME: syn_data.data_frame[LABEL_ID_COLUMN_NAME].values,
             }
         )
         if LABEL_ID_COLUMN_NAME in syn_data.data_frame.columns: