Update app.py

app.py

@@ -62,10 +62,20 @@ class Visualization:
         self.get_bboxes()
                 
     def get_cls_names(self):
-        with open(f"{self.root}/data.yaml", 'r') as file:
+        yaml_path = f"{self.root}/data.yaml"
+        if not os.path.exists(yaml_path):
+            print(f"Warning: {yaml_path} not found")
+            self.class_dict = {}
+            return
+            
+        with open(yaml_path, 'r') as file:
             data = yaml.safe_load(file)      
-        class_names = data['names']        
-        self.class_dict = {index: name for index, name in enumerate(class_names)}        
+        class_names = data.get('names', [])
+        self.class_dict = {index: name for index, name in enumerate(class_names)}
+        
+        # Print class names for debugging
+        if self.class_dict:
+            print(f"Dataset classes: {', '.join(class_names)}")        
     
     def get_bboxes(self):
         self.vis_datas, self.analysis_datas, self.im_paths = {}, {}, {}
@@ -291,8 +301,28 @@ def train_model(epochs, batch_size, img_size, device_selection):
         # Save the model path
         model_path = os.path.join(results_path, "weights", "best.pt")
         
+        # Load the trained model to ensure it's ready for inference
+        if os.path.exists(model_path):
+            model = YOLO(model_path)
+            class_info = f"\nTrained on {len(model.names)} classes: {', '.join(list(model.names.values())[:5])}"
+            if len(model.names) > 5:
+                class_info += f"... (총 {len(model.names)} 클래스)"
+        else:
+            class_info = ""
+        
         training_in_progress = False
-        return plots, f"Training completed! Model saved to {model_path}"
+        
+        # Provide instructions for saving the model
+        save_instructions = """
+        
+        ⚠️ **Important**: Models trained on Hugging Face Spaces are temporary!
+        To save your model permanently:
+        1. Download the model file: `./xray_detection/train/weights/best.pt`
+        2. Upload it to HuggingFace Hub or save locally
+        3. Use 'Load Pre-trained Model' with the saved path
+        """
+        
+        return plots, f"Training completed! Model saved to {model_path}{class_info}{save_instructions}"
         
     except Exception as e:
         training_in_progress = False
@@ -303,17 +333,32 @@ def run_inference(input_image, conf_threshold):
     """Run inference on a single image"""
     global model
     
+    # Try to load the trained model if not already loaded
     if model is None:
-        # Try to load a default model
-        try:
-            model = YOLO("yolo11n.pt")
-        except:
-            return None, "Please train the model first or load a pre-trained model!"
+        # First, try to load the trained model
+        trained_model_path = "./xray_detection/train/weights/best.pt"
+        if os.path.exists(trained_model_path):
+            try:
+                model = YOLO(trained_model_path)
+                print(f"Loaded trained model from {trained_model_path}")
+            except:
+                pass
+        
+        # If still no model, try to load default
+        if model is None:
+            try:
+                model = YOLO("yolo11n.pt")
+                print("Loaded default YOLOv11 model")
+            except:
+                return None, "Please train the model first or load a pre-trained model!"
     
     if input_image is None:
         return None, "Please upload an image!"
     
     try:
+        # Check if model is trained on X-ray dataset
+        model_info = f"Using model with {len(model.names)} classes\n"
+        
         # Save the input image temporarily
         temp_path = "temp_inference.jpg"
         input_image.save(temp_path)
@@ -326,20 +371,26 @@ def run_inference(input_image, conf_threshold):
         
         # Get detection info
         detections = []
-        if results[0].boxes is not None:
+        if results[0].boxes is not None and len(results[0].boxes) > 0:
             for box in results[0].boxes:
                 cls = int(box.cls)
                 conf = float(box.conf)
                 cls_name = model.names[cls]
                 detections.append(f"{cls_name}: {conf:.2f}")
+            
+            detection_text = model_info + "Detections:\n" + "\n".join(detections)
+        else:
+            # Check if it's because of wrong model
+            if len(model.names) == 80:  # COCO dataset has 80 classes
+                detection_text = model_info + "No objects detected.\n\n⚠️ Note: This appears to be a general COCO model. For X-ray baggage detection, please train the model on the X-ray dataset first."
+            else:
+                detection_text = model_info + "No objects detected at this confidence threshold.\nTry lowering the confidence threshold."
         
         # Clean up
         if os.path.exists(temp_path):
             os.remove(temp_path)
         
-        detection_text = "\n".join(detections) if detections else "No objects detected"
-        
-        return Image.fromarray(annotated_image), f"Detections:\n{detection_text}"
+        return Image.fromarray(annotated_image), detection_text
         
     except Exception as e:
         return None, f"Error during inference: {str(e)}"
@@ -349,11 +400,21 @@ def batch_inference(data_type, num_images):
     """Run inference on multiple images from test set"""
     global model
     
+    # Try to load the trained model if not already loaded
     if model is None:
-        try:
-            model = YOLO("yolo11n.pt")
-        except:
-            return [], "Please train the model first!"
+        trained_model_path = "./xray_detection/train/weights/best.pt"
+        if os.path.exists(trained_model_path):
+            try:
+                model = YOLO(trained_model_path)
+                print(f"Loaded trained model for batch inference")
+            except:
+                try:
+                    model = YOLO("yolo11n.pt")
+                    print("Loaded default model for batch inference")
+                except:
+                    return [], "Please train the model first!"
+        else:
+            return [], "No trained model found. Please train the model first!"
     
     if dataset_path is None:
         return [], "Please download the dataset first!"
@@ -375,15 +436,47 @@ def batch_inference(data_type, num_images):
             annotated = results[0].plot()
             results_images.append(Image.fromarray(annotated))
         
-        return results_images, f"Processed {len(results_images)} images from {data_type} dataset"
+        # Check model type
+        model_type = "X-ray detection model" if len(model.names) != 80 else "General COCO model"
+        
+        return results_images, f"Processed {len(results_images)} images from {data_type} dataset using {model_type}"
         
     except Exception as e:
         return [], f"Error during batch inference: {str(e)}"
 
-def load_pretrained_model(model_path):
+def get_dataset_info():
+    """Get information about the X-ray dataset classes"""
+    if dataset_path is None:
+        return "Dataset not downloaded yet."
+    
+    try:
+        yaml_path = f"{dataset_path}/data.yaml"
+        if not os.path.exists(yaml_path):
+            return "Dataset configuration file not found."
+        
+        with open(yaml_path, 'r') as file:
+            data = yaml.safe_load(file)
+        
+        class_names = data.get('names', [])
+        num_classes = len(class_names)
+        
+        info = f"**X-ray Baggage Dataset Info:**\n"
+        info += f"- Number of classes: {num_classes}\n"
+        info += f"- Classes: {', '.join(class_names)}\n"
+        info += f"\nThese are the prohibited items that the model will learn to detect."
+        
+        return info
+    except Exception as e:
+        return f"Error reading dataset info: {str(e)}"
     """Load a pre-trained model"""
     global model
     try:
+        # Check if it's a HuggingFace model path
+        if model_path.startswith("hf://") or "/" in model_path and not os.path.exists(model_path):
+            # Load from HuggingFace Hub
+            model = YOLO(model_path)
+            return f"Model loaded successfully from HuggingFace: {model_path}"
+        
         if not os.path.exists(model_path):
             # Try default paths
             default_paths = [
@@ -396,8 +489,17 @@ def load_pretrained_model(model_path):
                     model_path = path
                     break
         
-        model = YOLO(model_path)
-        return f"Model loaded successfully from {model_path}"
+        if os.path.exists(model_path):
+            model = YOLO(model_path)
+            # Check if it's a trained model by looking at class names
+            if hasattr(model, 'names') and len(model.names) > 0:
+                class_names = ", ".join([f"{i}: {name}" for i, name in model.names.items()][:5])
+                if len(model.names) > 5:
+                    class_names += f"... (총 {len(model.names)} 클래스)"
+                return f"Model loaded successfully from {model_path}\n클래스: {class_names}"
+            return f"Model loaded successfully from {model_path}"
+        else:
+            return "Model file not found. Please train a model first or provide a valid path."
     except Exception as e:
         return f"Error loading model: {str(e)}"
 
@@ -424,6 +526,12 @@ with gr.Blocks(title="X-ray Baggage Anomaly Detection", theme=gr.themes.Soft())
            ```
            KDATA_API={"username":"your_username","key":"your_api_key"}
            ```
+           
+        ### Model Persistence on Hugging Face Spaces
+        - Models trained on Spaces are **temporary** and will be lost when the Space restarts
+        - After training, download your model using the "Download Model" button
+        - For permanent storage, upload to HuggingFace Hub or save locally
+        - You can load saved models using the "Load Pre-trained Model" feature
         """)
     
     with gr.Tab("📊 Dataset"):
@@ -433,6 +541,13 @@ with gr.Blocks(title="X-ray Baggage Anomaly Detection", theme=gr.themes.Soft())
         
         download_btn.click(download_dataset, outputs=download_status)
         
+        # Dataset info section
+        with gr.Row():
+            dataset_info = gr.Markdown(value=get_dataset_info())
+            info_btn = gr.Button("🔄 Refresh Dataset Info", scale=0)
+        
+        info_btn.click(get_dataset_info, outputs=dataset_info)
+        
         gr.Markdown("### Visualize Dataset Samples")
         with gr.Row():
             data_type_viz = gr.Dropdown(["train", "valid", "test"], value="train", label="Dataset Type")
@@ -465,6 +580,8 @@ with gr.Blocks(title="X-ray Baggage Anomaly Detection", theme=gr.themes.Soft())
         - Use smaller batch sizes (4-8) to avoid GPU memory issues
         - Start with fewer epochs (5-10) for testing
         - Image size 480 provides good balance between quality and speed
+        
+        ⚠️ **Important**: Models are temporary on Spaces! Download your model after training.
         """)
         
         with gr.Row():
@@ -482,20 +599,68 @@ with gr.Blocks(title="X-ray Baggage Anomaly Detection", theme=gr.themes.Soft())
                        inputs=[epochs_input, batch_size_input, img_size_input, device_input],
                        outputs=[training_gallery, training_status])
         
-        gr.Markdown("### Load Pre-trained Model")
-        with gr.Row():
-            model_path_input = gr.Textbox(label="Model Path", value="./xray_detection/train/weights/best.pt")
-            load_model_btn = gr.Button("Load Model")
-            load_status = gr.Textbox(label="Status", interactive=False)
+        gr.Markdown("### Model Management")
         
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("#### Download Trained Model")
+                download_model_btn = gr.Button("Download Model (.pt)", variant="secondary")
+                download_file = gr.File(label="Download Model File", visible=False)
+                
+            with gr.Column():
+                gr.Markdown("#### Load Pre-trained Model")
+                model_path_input = gr.Textbox(
+                    label="Model Path", 
+                    value="./xray_detection/train/weights/best.pt",
+                    placeholder="Local path or HuggingFace model ID"
+                )
+                load_model_btn = gr.Button("Load Model")
+                load_status = gr.Textbox(label="Status", interactive=False)
+        
+        # Add download functionality
+        def download_trained_model():
+            model_path = "./xray_detection/train/weights/best.pt"
+            if os.path.exists(model_path):
+                return gr.update(value=model_path, visible=True)
+            else:
+                return gr.update(value=None, visible=True)
+        
+        download_model_btn.click(download_trained_model, outputs=download_file)
         load_model_btn.click(load_pretrained_model, inputs=model_path_input, outputs=load_status)
     
     with gr.Tab("🔍 Inference"):
+        # Model status check
+        def check_model_status():
+            global model
+            if model is None:
+                # Try to load trained model
+                trained_path = "./xray_detection/train/weights/best.pt"
+                if os.path.exists(trained_path):
+                    try:
+                        model = YOLO(trained_path)
+                        return f"✅ Trained model loaded: {len(model.names)} classes"
+                    except:
+                        return "❌ No model loaded. Please train or load a model first."
+                return "❌ No model loaded. Please train or load a model first."
+            else:
+                num_classes = len(model.names)
+                if num_classes == 80:
+                    return f"⚠️ Default COCO model loaded ({num_classes} classes). For X-ray detection, please train on the X-ray dataset."
+                else:
+                    return f"✅ Model loaded: {num_classes} classes - {', '.join(list(model.names.values())[:5])}..."
+        
+        with gr.Row():
+            model_status = gr.Textbox(label="Model Status", value=check_model_status(), interactive=False)
+            refresh_status_btn = gr.Button("🔄 Refresh Status", scale=0)
+        
+        refresh_status_btn.click(check_model_status, outputs=model_status)
+        
         gr.Markdown("### Single Image Inference")
+        gr.Markdown("Upload an X-ray baggage image to detect prohibited items.")
         
         with gr.Row():
             with gr.Column():
-                input_image = gr.Image(type="pil", label="Upload Image")
+                input_image = gr.Image(type="pil", label="Upload X-ray Image")
                 conf_threshold = gr.Slider(0.1, 0.9, 0.5, step=0.05, label="Confidence Threshold")
                 inference_btn = gr.Button("Run Detection", variant="primary")
             
@@ -508,6 +673,7 @@ with gr.Blocks(title="X-ray Baggage Anomaly Detection", theme=gr.themes.Soft())
                           outputs=[output_image, detection_info])
         
         gr.Markdown("### Batch Inference")
+        gr.Markdown("Run detection on multiple images from the test dataset.")
         
         with gr.Row():
             batch_data_type = gr.Dropdown(["test", "valid"], value="test", label="Dataset Type")