Create app.py

app.py

@@ -0,0 +1,407 @@
+import os
+import cv2
+import yaml
+import torch
+import random
+import gradio as gr
+import numpy as np
+import kagglehub
+from PIL import Image
+from glob import glob
+import matplotlib.pyplot as plt
+from matplotlib import patches
+from torchvision import transforms as T
+from ultralytics import YOLO
+import shutil
+import tempfile
+from pathlib import Path
+
+# 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", "")
+
+# Global variables
+model = None
+dataset_path = None
+training_in_progress = False
+
+class Visualization:
+    def __init__(self, root, data_types, n_ims, rows, cmap=None):
+        self.n_ims, self.rows = n_ims, rows
+        self.cmap, self.data_types = cmap, data_types
+        self.colors = ["firebrick", "darkorange", "blueviolet"]
+        self.root = root
+        
+        self.get_cls_names()
+        self.get_bboxes()
+                
+    def get_cls_names(self):
+        with open(f"{self.root}/data.yaml", 'r') as file:
+            data = yaml.safe_load(file)      
+        class_names = data['names']        
+        self.class_dict = {index: name for index, name in enumerate(class_names)}        
+    
+    def get_bboxes(self):
+        self.vis_datas, self.analysis_datas, self.im_paths = {}, {}, {}
+        for data_type in self.data_types:
+            all_bboxes, all_analysis_datas = [], {}
+            im_paths = glob(f"{self.root}/{data_type}/images/*")
+            
+            for idx, im_path in enumerate(im_paths):
+                bboxes = []
+                im_ext = os.path.splitext(im_path)[-1]
+                lbl_path = im_path.replace(im_ext, ".txt")
+                lbl_path = lbl_path.replace(f"{data_type}/images", f"{data_type}/labels")
+                if not os.path.isfile(lbl_path):
+                    continue                
+                meta_data = open(lbl_path).readlines()
+                for data in meta_data:                    
+                    parts = data.strip().split()[:5]
+                    cls_name = self.class_dict[int(parts[0])]
+                    bboxes.append([cls_name] + [float(x) for x in parts[1:]])
+                    if cls_name not in all_analysis_datas:
+                        all_analysis_datas[cls_name] = 1
+                    else:
+                        all_analysis_datas[cls_name] += 1
+                all_bboxes.append(bboxes)
+                    
+            self.vis_datas[data_type] = all_bboxes
+            self.analysis_datas[data_type] = all_analysis_datas
+            self.im_paths[data_type] = im_paths
+    
+    def plot_single(self, im_path, bboxes):
+        fig, ax = plt.subplots(figsize=(8, 8))
+        or_im = np.array(Image.open(im_path).convert("RGB"))
+        height, width, _ = or_im.shape
+
+        for bbox in bboxes:
+            class_id, x_center, y_center, w, h = bbox
+
+            x_min = int((x_center - w / 2) * width)
+            y_min = int((y_center - h / 2) * height)
+            x_max = int((x_center + w / 2) * width)
+            y_max = int((y_center + h / 2) * height)
+            
+            color = (random.randint(0, 255)/255, random.randint(0, 255)/255, random.randint(0, 255)/255)
+            cv2.rectangle(img=or_im, pt1=(x_min, y_min), pt2=(x_max, y_max), 
+                         color=(int(color[0]*255), int(color[1]*255), int(color[2]*255)), thickness=3)
+        
+        ax.imshow(or_im)
+        ax.axis("off")
+        ax.set_title(f"Number of objects: {len(bboxes)}")
+        
+        return fig
+
+    def vis_samples(self, data_type, n_samples=4):
+        if data_type not in self.vis_datas:
+            return None
+            
+        indices = [random.randint(0, len(self.vis_datas[data_type]) - 1) 
+                  for _ in range(min(n_samples, len(self.vis_datas[data_type])))]
+        
+        figs = []
+        for idx in indices:
+            im_path = self.im_paths[data_type][idx]
+            bboxes = self.vis_datas[data_type][idx]
+            fig = self.plot_single(im_path, bboxes)
+            figs.append(fig)
+            
+        return figs
+
+    def data_analysis(self, data_type):
+        if data_type not in self.analysis_datas:
+            return None
+            
+        plt.style.use('default')
+        fig, ax = plt.subplots(figsize=(12, 6))
+        
+        cls_names = list(self.analysis_datas[data_type].keys())
+        counts = list(self.analysis_datas[data_type].values())
+        
+        color_map = {"train": "firebrick", "valid": "darkorange", "test": "blueviolet"}
+        color = color_map.get(data_type, "steelblue")
+        
+        indices = np.arange(len(counts))
+        bars = ax.bar(indices, counts, 0.7, color=color)
+        
+        ax.set_xlabel("Class Names", fontsize=12)
+        ax.set_xticks(indices)
+        ax.set_xticklabels(cls_names, rotation=45, ha='right')
+        ax.set_ylabel("Data Counts", fontsize=12)
+        ax.set_title(f"{data_type.upper()} Dataset Class Distribution", fontsize=14)
+        
+        for i, (bar, v) in enumerate(zip(bars, counts)):
+            ax.text(bar.get_x() + bar.get_width()/2, v + 1, str(v), 
+                   ha='center', va='bottom', fontsize=10, color='navy')
+        
+        plt.tight_layout()
+        return fig
+
+def download_dataset():
+    """Download the dataset using kagglehub"""
+    global dataset_path
+    try:
+        dataset_path = kagglehub.dataset_download("orvile/x-ray-baggage-anomaly-detection")
+        return f"Dataset downloaded successfully to: {dataset_path}"
+    except Exception as e:
+        return f"Error downloading dataset: {str(e)}"
+
+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!"
+    
+    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)
+        return figs, f"Showing {len(figs)} samples from {data_type} dataset"
+    except Exception as e:
+        return None, f"Error visualizing data: {str(e)}"
+
+def analyze_class_distribution(data_type):
+    """Analyze class distribution in the dataset"""
+    if dataset_path is None:
+        return None, "Please download the dataset first!"
+    
+    try:
+        vis = Visualization(root=dataset_path, data_types=[data_type], 
+                          n_ims=20, rows=5, cmap="rgb")
+        fig = vis.data_analysis(data_type)
+        return fig, f"Class distribution for {data_type} dataset"
+    except Exception as e:
+        return None, f"Error analyzing data: {str(e)}"
+
+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!"
+    
+    if training_in_progress:
+        return None, "Training already in progress!"
+    
+    training_in_progress = True
+    
+    try:
+        # Determine device
+        if device_selection == "Auto":
+            device = 0 if torch.cuda.is_available() else "cpu"
+        elif device_selection == "CPU":
+            device = "cpu"
+        else:
+            device = 0
+        
+        # Initialize model
+        model = YOLO("yolo11n.pt")
+        
+        # Train model
+        results = model.train(
+            data=f"{dataset_path}/data.yaml",
+            epochs=epochs,
+            imgsz=img_size,
+            batch=batch_size,
+            device=device,
+            project="xray_detection",
+            name="train",
+            exist_ok=True,
+            verbose=True
+        )
+        
+        # Plot training results
+        results_path = "xray_detection/train"
+        plots = []
+        
+        for plot_file in ["results.png", "confusion_matrix.png", "val_batch0_pred.jpg"]:
+            plot_path = os.path.join(results_path, plot_file)
+            if os.path.exists(plot_path):
+                plots.append(Image.open(plot_path))
+        
+        training_in_progress = False
+        return plots, f"Training completed! Model saved to {results_path}"
+        
+    except Exception as e:
+        training_in_progress = False
+        return None, f"Error during training: {str(e)}"
+
+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:
+        # Save the input image temporarily
+        temp_path = "temp_inference.jpg"
+        input_image.save(temp_path)
+        
+        # Run inference
+        results = model(temp_path, conf=conf_threshold, verbose=False)
+        
+        # Draw results
+        annotated_image = results[0].plot()
+        
+        # Get detection info
+        detections = []
+        for r in results:
+            for box in r.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)
+        
+        detection_text = "\n".join(detections) if detections else "No objects detected"
+        
+        return Image.fromarray(annotated_image), f"Detections:\n{detection_text}"
+        
+    except Exception as e:
+        return None, f"Error during inference: {str(e)}"
+
+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!"
+    
+    if dataset_path is None:
+        return None, "Please download the dataset first!"
+    
+    try:
+        image_dir = f"{dataset_path}/{data_type}/images"
+        image_files = glob(f"{image_dir}/*")[:num_images]
+        
+        results_images = []
+        
+        for img_path in image_files:
+            results = model(img_path, verbose=False)
+            annotated = results[0].plot()
+            results_images.append(Image.fromarray(annotated))
+        
+        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)}"
+
+def load_pretrained_model(model_path):
+    """Load a pre-trained model"""
+    global model
+    try:
+        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:
+    gr.Markdown("""
+    # 🎯 X-ray Baggage Anomaly Detection with YOLOv11
+    
+    This application allows you to:
+    1. Download and visualize the X-ray baggage dataset
+    2. Analyze class distributions
+    3. Train a YOLOv11 model for object detection
+    4. Run inference on new images
+    """)
+    
+    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.click(download_dataset, outputs=download_status)
+        
+        gr.Markdown("### Visualize Dataset Samples")
+        with gr.Row():
+            data_type_viz = gr.Dropdown(["train", "valid", "test"], value="train", label="Dataset Type")
+            num_samples = gr.Slider(1, 8, 4, step=1, label="Number of Samples")
+            viz_btn = gr.Button("Visualize Samples")
+        
+        viz_gallery = gr.Gallery(label="Sample Images", columns=2, height="auto")
+        viz_status = gr.Textbox(label="Status", interactive=False)
+        
+        viz_btn.click(visualize_data, inputs=[data_type_viz, num_samples], 
+                     outputs=[viz_gallery, viz_status])
+        
+        gr.Markdown("### Analyze Class Distribution")
+        with gr.Row():
+            data_type_analysis = gr.Dropdown(["train", "valid", "test"], value="train", label="Dataset Type")
+            analyze_btn = gr.Button("Analyze Distribution")
+        
+        distribution_plot = gr.Plot(label="Class Distribution")
+        analysis_status = gr.Textbox(label="Status", interactive=False)
+        
+        analyze_btn.click(analyze_class_distribution, inputs=data_type_analysis, 
+                         outputs=[distribution_plot, analysis_status])
+    
+    with gr.Tab("🚀 Training"):
+        gr.Markdown("### Train YOLOv11 Model")
+        
+        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")
+            img_size_input = gr.Slider(320, 640, 480, step=32, label="Image Size")
+            device_input = gr.Radio(["Auto", "GPU", "CPU"], value="Auto", label="Device")
+        
+        train_btn = gr.Button("Start Training", variant="primary")
+        
+        training_gallery = gr.Gallery(label="Training Results", columns=3, height="auto")
+        training_status = gr.Textbox(label="Training Status", interactive=False)
+        
+        train_btn.click(train_model, 
+                       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="yolo11n.pt")
+            load_model_btn = gr.Button("Load Model")
+            load_status = gr.Textbox(label="Status", interactive=False)
+        
+        load_model_btn.click(load_pretrained_model, inputs=model_path_input, outputs=load_status)
+    
+    with gr.Tab("🔍 Inference"):
+        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)
+        
+        inference_btn.click(run_inference, 
+                          inputs=[input_image, conf_threshold],
+                          outputs=[output_image, detection_info])
+        
+        gr.Markdown("### Batch Inference")
+        
+        with gr.Row():
+            batch_data_type = gr.Dropdown(["test", "valid"], value="test", label="Dataset Type")
+            batch_num_images = gr.Slider(1, 10, 5, step=1, label="Number of Images")
+            batch_btn = gr.Button("Run Batch Inference")
+        
+        batch_gallery = gr.Gallery(label="Batch Results", columns=3, height="auto")
+        batch_status = gr.Textbox(label="Status", interactive=False)
+        
+        batch_btn.click(batch_inference,
+                       inputs=[batch_data_type, batch_num_images],
+                       outputs=[batch_gallery, batch_status])
+
+# Launch the app
+if __name__ == "__main__":
+    demo.launch(share=True)