detecton endpoints

app.py

@@ -250,7 +250,6 @@
 # if __name__ == "__main__":
 #     demo.launch()
 
-
 import gradio as gr
 import cv2
 import numpy as np
@@ -360,20 +359,24 @@ class EnginePartDetector:
             try:
                 with open(self.TEMPLATE_FILE, "rb") as f:
                     data = pickle.load(f)
-                    # Support legacy format if needed, but here we assume the new format
-                    if isinstance(data, dict):
-                        # If old format was {name: {"features": feat, "roi": roi}}
-                        # we convert it to {name: [feat]}
+                    
+                if isinstance(data, dict):
+                    # Check if it's the new combined format
+                    if "classes" in data and "rois" in data:
+                        self.classes = data["classes"]
+                        self.class_rois = data["rois"]
+                    # Check if it's the temporary refactored format
+                    elif all(isinstance(v, list) for v in data.values()):
+                        self.classes = data
+                        self.class_rois = {}
+                    # Legacy support: {name: {"features": feat, "roi": roi}}
+                    else:
                         self.classes = {}
                         self.class_rois = {}
                         for k, v in data.items():
                             if isinstance(v, dict) and "features" in v:
                                 self.classes[k] = [v["features"]]
                                 self.class_rois[k] = v.get("roi")
-                            else:
-                                self.classes[k] = v
-                    else:
-                        self.classes = {}
                 logger.info(f"Loaded {len(self.classes)} class(es).")
             except Exception as e:
                 logger.error(f"Data load failed: {e}")
@@ -381,12 +384,12 @@ class EnginePartDetector:
 
     def _persist_data(self) -> None:
         try:
+            # Save everything in one file to ensure ROIs are preserved
             with open(self.TEMPLATE_FILE, "wb") as f:
-                pickle.dump(self.classes, f)
-            # Separately save ROIs if needed, but for now we just persist classes
-            # In a real app we'd save ROIs too. Let's include them in a combined dict.
-            with open("class_data.pkl", "wb") as f:
-                pickle.dump({"classes": self.classes, "rois": self.class_rois}, f)
+                pickle.dump({
+                    "classes": self.classes,
+                    "rois": self.class_rois
+                }, f)
         except Exception as e:
             logger.error(f"Data save failed: {e}")
 
@@ -628,7 +631,7 @@ custom_css = """
 .footer { text-align: center; margin-top: 2rem; color: #666; }
 """
 
-with gr.Blocks(title="Engine Part CV System", theme=gr.themes.Soft(), css=custom_css) as demo:
+with gr.Blocks(title="Engine Part CV System") as demo:
     gr.Markdown("""
     <div class="header">
         <h1>🔧 Engine Part CV System</h1>
@@ -694,7 +697,7 @@ with gr.Blocks(title="Engine Part CV System", theme=gr.themes.Soft(), css=custom
                 return detector.list_templates(), detector.get_template_roi(first_name)
             return "No classes trained yet.", None
 
-        refresh_btn.click(fn=update_library_preview, outputs=[template_list, library_roi_view])
+        refresh_btn.click(fn=update_library_preview, outputs=[template_list, library_roi_view], api_name="list_classes")
         demo.load(fn=update_library_preview, outputs=[template_list, library_roi_view])
 
     gr.Markdown("""
@@ -705,4 +708,9 @@ with gr.Blocks(title="Engine Part CV System", theme=gr.themes.Soft(), css=custom
     """)
 
 if __name__ == "__main__":
-    demo.launch(share=False, show_error=True)
+    demo.launch(
+        share=False, 
+        show_error=True,
+        theme=gr.themes.Soft(),
+        css=custom_css
+    )