Update app.py

app.py

@@ -125,10 +125,8 @@ class Visualization:
         cv2.putText(or_im, f"Objects: {len(bboxes)}", (10, 30), 
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
         
-        # Convert BGR to RGB if needed
-        if len(or_im.shape) == 3 and or_im.shape[2] == 3:
-            or_im = cv2.cvtColor(or_im, cv2.COLOR_BGR2RGB)
-        
+        # OpenCV uses BGR, but PIL expects RGB, and we already loaded as RGB
+        # So no conversion needed
         return Image.fromarray(or_im)
 
     def vis_samples(self, data_type, n_samples=4):
@@ -259,8 +257,12 @@ def train_model(epochs, batch_size, img_size, device_selection):
         else:
             device = 0 if torch.cuda.is_available() else "cpu"
         
-        # Initialize model
-        model = YOLO("yolo11n.pt")
+        # Initialize model - use yolov8n if yolo11n not available
+        try:
+            model = YOLO("yolo11n.pt")
+        except Exception as e:
+            print(f"YOLOv11 not available: {e}, falling back to YOLOv8")
+            model = YOLO("yolov8n.pt")  # Fallback to YOLOv8
         
         # Create project directory
         project_dir = "./xray_detection"
@@ -350,14 +352,23 @@ def run_inference(input_image, conf_threshold):
                 model = YOLO("yolo11n.pt")
                 print("Loaded default YOLOv11 model")
             except:
-                return None, "Please train the model first or load a pre-trained model!"
+                try:
+                    model = YOLO("yolov8n.pt")
+                    print("Loaded YOLOv8 model as fallback")
+                except Exception as e:
+                    print(f"Failed to load default model: {e}")
+                    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"
+        model_info = ""
+        try:
+            model_info = f"Using model with {len(model.names)} classes\n"
+        except:
+            model_info = "Using loaded model\n"
         
         # Save the input image temporarily
         temp_path = "temp_inference.jpg"
@@ -375,16 +386,22 @@ def run_inference(input_image, conf_threshold):
             for box in results[0].boxes:
                 cls = int(box.cls)
                 conf = float(box.conf)
-                cls_name = model.names[cls]
+                try:
+                    cls_name = model.names[cls]
+                except:
+                    cls_name = f"Class {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."
+            try:
+                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."
+            except:
+                detection_text = model_info + "No objects detected."
         
         # Clean up
         if os.path.exists(temp_path):
@@ -412,7 +429,11 @@ def batch_inference(data_type, num_images):
                     model = YOLO("yolo11n.pt")
                     print("Loaded default model for batch inference")
                 except:
-                    return [], "Please train the model first!"
+                    try:
+                        model = YOLO("yolov8n.pt")
+                        print("Loaded YOLOv8 model as fallback for batch inference")
+                    except:
+                        return [], "Please train the model first!"
         else:
             return [], "No trained model found. Please train the model first!"
     
@@ -437,7 +458,10 @@ def batch_inference(data_type, num_images):
             results_images.append(Image.fromarray(annotated))
         
         # Check model type
-        model_type = "X-ray detection model" if len(model.names) != 80 else "General COCO model"
+        try:
+            model_type = "X-ray detection model" if len(model.names) != 80 else "General COCO model"
+        except:
+            model_type = "Loaded model"
         
         return results_images, f"Processed {len(results_images)} images from {data_type} dataset using {model_type}"
         
@@ -482,7 +506,8 @@ def get_dataset_info():
             default_paths = [
                 "./xray_detection/train/weights/best.pt",
                 "./xray_detection/train/weights/last.pt",
-                "yolo11n.pt"
+                "yolo11n.pt",
+                "yolov8n.pt"
             ]
             for path in default_paths:
                 if os.path.exists(path):
@@ -492,26 +517,55 @@ def get_dataset_info():
         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}"
+            try:
+                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}"
+            except:
+                pass
             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)}"
 
+def check_model_status():
+    """Check current 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:
+        try:
+            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:
+                class_names = ', '.join(list(model.names.values())[:5])
+                if len(model.names) > 5:
+                    class_names += "..."
+                return f"✅ Model loaded: {num_classes} classes - {class_names}"
+        except:
+            return "✅ Model loaded"
+
 # Create Gradio interface
 with gr.Blocks(title="X-ray Baggage Anomaly Detection", theme=gr.themes.Soft()) as demo:
     gr.Markdown("""
-    # 🎯 X-ray Baggage Anomaly Detection with YOLOv11
+    # 🎯 X-ray Baggage Anomaly Detection with YOLO
     
     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
+    3. Train a YOLO model for object detection
     4. Run inference on new images
     
     **Note:** GPU will be automatically allocated when needed for training and inference.
@@ -543,10 +597,36 @@ with gr.Blocks(title="X-ray Baggage Anomaly Detection", theme=gr.themes.Soft())
         
         # Dataset info section
         with gr.Row():
-            dataset_info = gr.Markdown(value=get_dataset_info())
+            dataset_info = gr.Markdown(value="Dataset not downloaded yet.")
             info_btn = gr.Button("🔄 Refresh Dataset Info", scale=0)
         
-        info_btn.click(get_dataset_info, outputs=dataset_info)
+        def update_dataset_info():
+            return get_dataset_info()
+        
+        info_btn.click(update_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")
+            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.Image(label="Class Distribution", type="pil")
+        analysis_status = gr.Textbox(label="Status", interactive=False)
+        
+        analyze_btn.click(analyze_class_distribution, inputs=data_type_analysis, 
+                         outputs=[distribution_plot, analysis_status])
         
         gr.Markdown("### Visualize Dataset Samples")
         with gr.Row():
@@ -572,7 +652,7 @@ with gr.Blocks(title="X-ray Baggage Anomaly Detection", theme=gr.themes.Soft())
                          outputs=[distribution_plot, analysis_status])
     
     with gr.Tab("🚀 Training"):
-        gr.Markdown("### Train YOLOv11 Model")
+        gr.Markdown("### Train YOLO Model")
         gr.Markdown("""
         **Note:** Training will automatically use GPU if available. This may take several minutes.
         
@@ -601,6 +681,14 @@ with gr.Blocks(title="X-ray Baggage Anomaly Detection", theme=gr.themes.Soft())
         
         gr.Markdown("### Model Management")
         
+        # Define download_trained_model function here
+        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)
+        
         with gr.Row():
             with gr.Column():
                 gr.Markdown("#### Download Trained Model")
@@ -617,38 +705,11 @@ with gr.Blocks(title="X-ray Baggage Anomaly Detection", theme=gr.themes.Soft())
                 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)