Update app.py

app.py

@@ -15,13 +15,27 @@ from ultralytics import YOLO
 import shutil
 import tempfile
 from pathlib import Path
+import json
+
+# Try to import spaces for Hugging Face Spaces GPU support
+try:
+    import spaces
+    ON_SPACES = True
+except ImportError:
+    ON_SPACES = False
+    # Create a dummy decorator if not on Spaces
+    class spaces:
+        @staticmethod
+        def GPU(duration=60):
+            def decorator(func):
+                return func
+            return decorator
 
 # Set Kaggle API credentials from environment variable
 if os.getenv("KDATA_API"):
     kaggle_key = os.getenv("KDATA_API")
     # Parse the key if it's in JSON format
     if "{" in kaggle_key:
-        import json
         key_data = json.loads(kaggle_key)
         os.environ["KAGGLE_USERNAME"] = key_data.get("username", "")
         os.environ["KAGGLE_KEY"] = key_data.get("key", "")
@@ -95,6 +109,7 @@ class Visualization:
         ax.imshow(or_im)
         ax.axis("off")
         ax.set_title(f"Number of objects: {len(bboxes)}")
+        plt.tight_layout()
         
         return fig
 
@@ -147,23 +162,37 @@ def download_dataset():
     """Download the dataset using kagglehub"""
     global dataset_path
     try:
+        # Create a local directory to store the dataset
+        local_dir = "./xray_dataset"
+        
+        # Download dataset
         dataset_path = kagglehub.dataset_download("orvile/x-ray-baggage-anomaly-detection")
+        
+        # If the dataset is downloaded to a temporary location, copy it to our local directory
+        if dataset_path != local_dir and os.path.exists(dataset_path):
+            if os.path.exists(local_dir):
+                shutil.rmtree(local_dir)
+            shutil.copytree(dataset_path, local_dir)
+            dataset_path = local_dir
+        
         return f"Dataset downloaded successfully to: {dataset_path}"
     except Exception as e:
-        return f"Error downloading dataset: {str(e)}"
+        return f"Error downloading dataset: {str(e)}\n\nPlease ensure KDATA_API environment variable is set correctly."
 
 def visualize_data(data_type, num_samples):
     """Visualize sample images from the dataset"""
     if dataset_path is None:
-        return None, "Please download the dataset first!"
+        return [], "Please download the dataset first!"
     
     try:
         vis = Visualization(root=dataset_path, data_types=[data_type], 
                           n_ims=num_samples, rows=2, cmap="rgb")
         figs = vis.vis_samples(data_type, num_samples)
+        if figs is None:
+            return [], f"No data found for {data_type} dataset"
         return figs, f"Showing {len(figs)} samples from {data_type} dataset"
     except Exception as e:
-        return None, f"Error visualizing data: {str(e)}"
+        return [], f"Error visualizing data: {str(e)}"
 
 def analyze_class_distribution(data_type):
     """Analyze class distribution in the dataset"""
@@ -174,34 +203,43 @@ def analyze_class_distribution(data_type):
         vis = Visualization(root=dataset_path, data_types=[data_type], 
                           n_ims=20, rows=5, cmap="rgb")
         fig = vis.data_analysis(data_type)
+        if fig is None:
+            return None, f"No data found for {data_type} dataset"
         return fig, f"Class distribution for {data_type} dataset"
     except Exception as e:
         return None, f"Error analyzing data: {str(e)}"
 
+@spaces.GPU(duration=300)  # Request GPU for 5 minutes for training
 def train_model(epochs, batch_size, img_size, device_selection):
     """Train YOLOv11 model"""
     global model, training_in_progress
     
     if dataset_path is None:
-        return None, "Please download the dataset first!"
+        return [], "Please download the dataset first!"
     
     if training_in_progress:
-        return None, "Training already in progress!"
+        return [], "Training already in progress!"
     
     training_in_progress = True
     
     try:
-        # Determine device
-        if device_selection == "Auto":
+        # Determine device - on Spaces, always use GPU if available
+        if ON_SPACES and torch.cuda.is_available():
+            device = 0
+        elif device_selection == "Auto":
             device = 0 if torch.cuda.is_available() else "cpu"
         elif device_selection == "CPU":
             device = "cpu"
         else:
-            device = 0
+            device = 0 if torch.cuda.is_available() else "cpu"
         
         # Initialize model
         model = YOLO("yolo11n.pt")
         
+        # Create project directory
+        project_dir = "./xray_detection"
+        os.makedirs(project_dir, exist_ok=True)
+        
         # Train model
         results = model.train(
             data=f"{dataset_path}/data.yaml",
@@ -209,34 +247,50 @@ def train_model(epochs, batch_size, img_size, device_selection):
             imgsz=img_size,
             batch=batch_size,
             device=device,
-            project="xray_detection",
+            project=project_dir,
             name="train",
             exist_ok=True,
-            verbose=True
+            verbose=True,
+            patience=5,  # Reduce patience for faster training on Spaces
+            save_period=5  # Save checkpoints every 5 epochs
         )
         
-        # Plot training results
-        results_path = "xray_detection/train"
+        # Collect training result plots
+        results_path = os.path.join(project_dir, "train")
         plots = []
         
-        for plot_file in ["results.png", "confusion_matrix.png", "val_batch0_pred.jpg"]:
+        plot_files = ["results.png", "confusion_matrix.png", "val_batch0_pred.jpg", 
+                     "train_batch0.jpg", "val_batch0_labels.jpg"]
+        
+        for plot_file in plot_files:
             plot_path = os.path.join(results_path, plot_file)
             if os.path.exists(plot_path):
                 plots.append(Image.open(plot_path))
         
+        # Save the model path
+        model_path = os.path.join(results_path, "weights", "best.pt")
+        
         training_in_progress = False
-        return plots, f"Training completed! Model saved to {results_path}"
+        return plots, f"Training completed! Model saved to {model_path}"
         
     except Exception as e:
         training_in_progress = False
-        return None, f"Error during training: {str(e)}"
+        return [], f"Error during training: {str(e)}"
 
+@spaces.GPU(duration=60)  # Request GPU for 1 minute for inference
 def run_inference(input_image, conf_threshold):
     """Run inference on a single image"""
     global model
     
     if model is None:
-        return None, "Please train the model first or load a pre-trained model!"
+        # 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!"
+    
+    if input_image is None:
+        return None, "Please upload an image!"
     
     try:
         # Save the input image temporarily
@@ -251,15 +305,16 @@ def run_inference(input_image, conf_threshold):
         
         # Get detection info
         detections = []
-        for r in results:
-            for box in r.boxes:
+        if results[0].boxes is not None:
+            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}")
         
         # Clean up
-        os.remove(temp_path)
+        if os.path.exists(temp_path):
+            os.remove(temp_path)
         
         detection_text = "\n".join(detections) if detections else "No objects detected"
         
@@ -268,20 +323,30 @@ def run_inference(input_image, conf_threshold):
     except Exception as e:
         return None, f"Error during inference: {str(e)}"
 
+@spaces.GPU(duration=60)  # Request GPU for batch inference
 def batch_inference(data_type, num_images):
     """Run inference on multiple images from test set"""
     global model
     
     if model is None:
-        return None, "Please train the model first!"
+        try:
+            model = YOLO("yolo11n.pt")
+        except:
+            return [], "Please train the model first!"
     
     if dataset_path is None:
-        return None, "Please download the dataset first!"
+        return [], "Please download the dataset first!"
     
     try:
         image_dir = f"{dataset_path}/{data_type}/images"
+        if not os.path.exists(image_dir):
+            return [], f"Directory {image_dir} not found!"
+            
         image_files = glob(f"{image_dir}/*")[:num_images]
         
+        if not image_files:
+            return [], f"No images found in {image_dir}"
+        
         results_images = []
         
         for img_path in image_files:
@@ -292,19 +357,31 @@ def batch_inference(data_type, num_images):
         return results_images, f"Processed {len(results_images)} images from {data_type} dataset"
         
     except Exception as e:
-        return None, f"Error during batch inference: {str(e)}"
+        return [], f"Error during batch inference: {str(e)}"
 
 def load_pretrained_model(model_path):
     """Load a pre-trained model"""
     global model
     try:
+        if not os.path.exists(model_path):
+            # Try default paths
+            default_paths = [
+                "./xray_detection/train/weights/best.pt",
+                "./xray_detection/train/weights/last.pt",
+                "yolo11n.pt"
+            ]
+            for path in default_paths:
+                if os.path.exists(path):
+                    model_path = path
+                    break
+        
         model = YOLO(model_path)
         return f"Model loaded successfully from {model_path}"
     except Exception as e:
         return f"Error loading model: {str(e)}"
 
 # Create Gradio interface
-with gr.Blocks(title="X-ray Baggage Anomaly Detection") as demo:
+with gr.Blocks(title="X-ray Baggage Anomaly Detection", theme=gr.themes.Soft()) as demo:
     gr.Markdown("""
     # 🎯 X-ray Baggage Anomaly Detection with YOLOv11
     
@@ -313,12 +390,25 @@ with gr.Blocks(title="X-ray Baggage Anomaly Detection") as demo:
     2. Analyze class distributions
     3. Train a YOLOv11 model for object detection
     4. Run inference on new images
+    
+    **Note:** GPU will be automatically allocated when needed for training and inference.
     """)
     
+    # Add instructions for Kaggle API setup
+    with gr.Accordion("📝 Setup Instructions", open=False):
+        gr.Markdown("""
+        ### Kaggle API Setup
+        1. Get your Kaggle API credentials from https://www.kaggle.com/settings
+        2. Set the KDATA_API environment variable in Hugging Face Spaces settings:
+           ```
+           KDATA_API={"username":"your_username","key":"your_api_key"}
+           ```
+        """)
+    
     with gr.Tab("📊 Dataset"):
         with gr.Row():
-            download_btn = gr.Button("Download Dataset", variant="primary")
-            download_status = gr.Textbox(label="Status", interactive=False)
+            download_btn = gr.Button("Download Dataset", variant="primary", scale=1)
+            download_status = gr.Textbox(label="Status", interactive=False, scale=3)
         
         download_btn.click(download_dataset, outputs=download_status)
         
@@ -347,6 +437,7 @@ with gr.Blocks(title="X-ray Baggage Anomaly Detection") as demo:
     
     with gr.Tab("🚀 Training"):
         gr.Markdown("### Train YOLOv11 Model")
+        gr.Markdown("**Note:** Training will automatically use GPU if available. This may take several minutes.")
         
         with gr.Row():
             epochs_input = gr.Slider(1, 50, 10, step=1, label="Epochs")
@@ -365,7 +456,7 @@ with gr.Blocks(title="X-ray Baggage Anomaly Detection") as demo:
         
         gr.Markdown("### Load Pre-trained Model")
         with gr.Row():
-            model_path_input = gr.Textbox(label="Model Path", value="yolo11n.pt")
+            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)
         
@@ -375,14 +466,14 @@ with gr.Blocks(title="X-ray Baggage Anomaly Detection") as demo:
         gr.Markdown("### Single Image Inference")
         
         with gr.Row():
-            input_image = gr.Image(type="pil", label="Upload 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")
-        
-        with gr.Row():
-            output_image = gr.Image(type="pil", label="Detection Result")
-            detection_info = gr.Textbox(label="Detection Info", lines=5)
+            with gr.Column():
+                input_image = gr.Image(type="pil", label="Upload 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")
+            
+            with gr.Column():
+                output_image = gr.Image(type="pil", label="Detection Result")
+                detection_info = gr.Textbox(label="Detection Info", lines=5)
         
         inference_btn.click(run_inference, 
                           inputs=[input_image, conf_threshold],
@@ -404,4 +495,8 @@ with gr.Blocks(title="X-ray Baggage Anomaly Detection") as demo:
 
 # Launch the app
 if __name__ == "__main__":
-    demo.launch(share=True)
+    # Check if running on Hugging Face Spaces
+    if ON_SPACES:
+        demo.launch(ssr_mode=False)
+    else:
+        demo.launch(share=True, ssr_mode=False)