fix: fix bug of image switch

app.py

@@ -62,12 +62,8 @@ class GUI:
     def get_avatar(self, idx):
         if idx == -1:
             img = cv2.imread("./imgs/db_diag/-1.png")
-        elif self.task == "prisoner_dilemma":
+        elif self.task == "simulation/prisoner_dilemma":
             img = cv2.imread(f"./imgs/prison/{idx}.png")
-        elif self.task == "db_diag":
-            img = cv2.imread(f"./imgs/db_diag/{idx}.png")
-        elif "sde" in self.task:
-            img = cv2.imread(f"./imgs/sde/{idx}.png")
         else:
             img = cv2.imread(f"./imgs/{idx}.png")
         base64_str = cv2.imencode(".png", img)[1].tostring()
@@ -178,6 +174,13 @@ class GUI:
             self.backend = TaskSolving.from_task(task_dropdown, self.tasks_dir)
         else:
             self.backend = Simulation.from_task(task_dropdown, self.tasks_dir)
+        self.agent_id = {
+            self.backend.agents[idx].name: idx
+            for idx in range(len(self.backend.agents))
+        }
+        
+        self.task = task_dropdown
+        self.stu_num = len(self.agent_id) - 1
         self.backend.reset()
         self.turns_remain = self.backend.environment.max_turns
 
@@ -218,7 +221,7 @@ class GUI:
         # if len(data) != self.stu_num:
         if len(data) != self.stu_num + 1:
             raise gr.Error("data length is not equal to the total number of students.")
-        if self.task == "prisoner_dilemma":
+        if self.task == "simulation/prisoner_dilemma":
             img = cv2.imread("./imgs/speaking.png", cv2.IMREAD_UNCHANGED)
             if (
                 len(self.messages) < 2
@@ -327,7 +330,9 @@ class GUI:
         """
 
         # data = self.backend.next_data()
+
         return_message = self.backend.next()
+
         data = self.return_format(return_message)
 
         # data.sort(key=lambda item: item["sender"])
@@ -460,7 +465,7 @@ class GUI:
                     task_dropdown = gr.Dropdown(
                         choices=[
                             "simulation/nlp_classroom_9players",
-                            "simulation/prisoner_dilemma"
+                            "simulation/prisoner_dilemma",
                         ],
                         value="simulation/nlp_classroom_9players",
                         label="Task",
@@ -590,6 +595,5 @@ class GUI:
         demo.queue(concurrency_count=5, max_size=20).launch()
         # demo.launch()
 
-
         
 GUI().launch()