Update app.py

app.py

@@ -1,4 +1,7 @@
 import os
+# Set environment variables for Spaces compatibility
+os.environ['OMP_NUM_THREADS'] = '1'
+os.environ['MKL_NUM_THREADS'] = '1'
 import cv2
 import yaml
 import torch
@@ -253,7 +256,7 @@ def train_model(epochs, batch_size, img_size, device_selection):
         project_dir = "./xray_detection"
         os.makedirs(project_dir, exist_ok=True)
         
-        # Train model
+        # Train model with workers=0 to avoid multiprocessing issues on Spaces
         results = model.train(
             data=f"{dataset_path}/data.yaml",
             epochs=epochs,
@@ -265,7 +268,12 @@ def train_model(epochs, batch_size, img_size, device_selection):
             exist_ok=True,
             verbose=True,
             patience=5,  # Reduce patience for faster training on Spaces
-            save_period=5  # Save checkpoints every 5 epochs
+            save_period=5,  # Save checkpoints every 5 epochs
+            workers=0,  # Important: Set to 0 to avoid multiprocessing issues
+            single_cls=False,
+            rect=False,
+            cache=False,  # Disable caching to avoid memory issues
+            amp=True  # Use automatic mixed precision for faster training
         )
         
         # Collect training result plots
@@ -310,8 +318,8 @@ def run_inference(input_image, conf_threshold):
         temp_path = "temp_inference.jpg"
         input_image.save(temp_path)
         
-        # Run inference
-        results = model(temp_path, conf=conf_threshold, verbose=False)
+        # Run inference with workers=0
+        results = model(temp_path, conf=conf_threshold, verbose=False, device=0 if torch.cuda.is_available() else 'cpu')
         
         # Draw results
         annotated_image = results[0].plot()
@@ -450,11 +458,18 @@ with gr.Blocks(title="X-ray Baggage Anomaly Detection", theme=gr.themes.Soft())
     
     with gr.Tab("🚀 Training"):
         gr.Markdown("### Train YOLOv11 Model")
-        gr.Markdown("**Note:** Training will automatically use GPU if available. This may take several minutes.")
+        gr.Markdown("""
+        **Note:** Training will automatically use GPU if available. This may take several minutes.
+        
+        **Tips for Hugging Face Spaces:**
+        - 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
+        """)
         
         with gr.Row():
             epochs_input = gr.Slider(1, 50, 10, step=1, label="Epochs")
-            batch_size_input = gr.Slider(8, 64, 16, step=8, label="Batch Size")
+            batch_size_input = gr.Slider(4, 32, 8, step=4, label="Batch Size (lower for limited GPU)")
             img_size_input = gr.Slider(320, 640, 480, step=32, label="Image Size")
             device_input = gr.Radio(["Auto", "GPU", "CPU"], value="Auto", label="Device")