reconstruc baseline, mask

app.py

@@ -7,6 +7,8 @@ import gradio as gr
 import spaces
 from pnpxai.core.experiment.auto_explanation import AutoExplanationForImageClassification
 from pnpxai.core.detector.detector import extract_graph_data, symbolic_trace
+from pnpxai.explainers.utils.baselines import BASELINE_FUNCTIONS_FOR_IMAGE
+from pnpxai.explainers.utils.feature_masks import FEATURE_MASK_FUNCTIONS_FOR_IMAGE
 import matplotlib.pyplot as plt
 import plotly.graph_objects as go
 import plotly.express as px
@@ -541,10 +543,9 @@ class ExplainerCheckbox(Component):
                         break
 
                 opt_exp_id = max([x['id'] for x in checkbox_group_info]) + 1
-                # opt_output.explainer.model = self.experiment.model
-                # self.experiment.manager._explainers.append(opt_output.explainer)
-                # self.experiment.manager._explainer_ids.append(opt_exp_id)
 
+                # Deliver the parameter and class and reconstruct
+                # It should be done because spaces.GPU cannot pickle the class object
                 opt_res = {
                     'id': opt_exp_id,
                     'class': opt_output.explainer.__class__,
@@ -558,15 +559,56 @@ class ExplainerCheckbox(Component):
                 return [opt_res, checkbox_group_info, checkbox, bttn]
 
             def update_exp(exp_res):
-                kwargs = {}
-                for k,v in exp_res['params'].items():
-                    if "explainer" in k:
-                        _key = k.split("explainer.")[1]
-                        kwargs[_key] = v
-                
-                kwargs['model'] = self.experiment.model
-                explainer = exp_res['class'](**kwargs)
-                _id = exp_res['id']
+                try:
+                    kwargs = {}
+                    has_baseline = False
+                    has_feature_mask = False
+                    for k,v in exp_res['params'].items():
+                        if "explainer" in k:
+                            _key = k.split("explainer.")[1]
+                            kwargs[_key] = v
+                            if "baseline_fn" in _key:
+                                has_baseline = True
+                            if "feature_mask_fn" in _key:
+                                has_feature_mask = True
+
+                    # Reconstruct baseline object
+                    if has_baseline:
+                        method = kwargs['baseline_fn.method']
+                        del kwargs['baseline_fn.method']
+                        baseline_kwargs = {}
+                        keys = list(kwargs.keys())
+                        for k in keys:
+                            v = kwargs[k]
+                            if "baseline_fn" in k:
+                                baseline_kwargs[k.split("baseline_fn.")[1]] = v
+                                del kwargs[k]
+                        if method == "mean":
+                            baseline_kwargs['dim'] = 1 # Set arbitrary value
+                        baseline_fn = BASELINE_FUNCTIONS_FOR_IMAGE[method](**baseline_kwargs)
+                        kwargs['baseline_fn'] = baseline_fn
+
+                    # Reconstruct feature_mask object
+                    if has_feature_mask:
+                        method = kwargs['feature_mask_fn.method']
+                        del kwargs['feature_mask_fn.method']
+                        mask_kwargs = {}
+                        keys = list(kwargs.keys())
+                        for k in keys:
+                            v = kwargs[k]
+                            if "feature_mask_fn" in k:
+                                mask_kwargs[k.split("feature_mask_fn.")[1]] = v
+                                del kwargs[k]
+                        mask_fn = FEATURE_MASK_FUNCTIONS_FOR_IMAGE[method](**mask_kwargs)
+                        kwargs['feature_mask_fn'] = mask_fn
+                    
+                    kwargs['model'] = self.experiment.model
+                    explainer = exp_res['class'](**kwargs)
+                    _id = exp_res['id']
+                except Exception as e:
+                    # If the optimization is failed, use the default parameter explainer as optimal
+                    explainer = self.experiment.manager._explainers[self.default_exp_id]
+
                 self.experiment.manager._explainers.append(explainer)
                 self.experiment.manager._explainer_ids.append(_id)