Update app.py

app.py

@@ -257,6 +257,14 @@ def train_model(epochs, batch_size, img_size, device_selection):
         else:
             device = 0 if torch.cuda.is_available() else "cpu"
         
+        # Read dataset info
+        yaml_path = f"{dataset_path}/data.yaml"
+        with open(yaml_path, 'r') as file:
+            data_config = yaml.safe_load(file)
+        
+        class_names = data_config.get('names', [])
+        print(f"Training on {len(class_names)} classes: {class_names}")
+        
         # Initialize model - use yolov8n if yolo11n not available
         try:
             model = YOLO("yolo11n.pt")
@@ -268,9 +276,9 @@ def train_model(epochs, batch_size, img_size, device_selection):
         project_dir = "./xray_detection"
         os.makedirs(project_dir, exist_ok=True)
         
-        # Train model with workers=0 to avoid multiprocessing issues on Spaces
+        # Train model with optimized settings for X-ray detection
         results = model.train(
-            data=f"{dataset_path}/data.yaml",
+            data=yaml_path,
             epochs=epochs,
             imgsz=img_size,
             batch=batch_size,
@@ -285,7 +293,34 @@ def train_model(epochs, batch_size, img_size, device_selection):
             single_cls=False,
             rect=False,
             cache=False,  # Disable caching to avoid memory issues
-            amp=True  # Use automatic mixed precision for faster training
+            amp=True,  # Use automatic mixed precision for faster training
+            # Optimization settings
+            optimizer='AdamW',
+            lr0=0.001,  # Initial learning rate
+            lrf=0.01,  # Final learning rate factor
+            momentum=0.937,
+            weight_decay=0.0005,
+            warmup_epochs=3.0,
+            warmup_momentum=0.8,
+            warmup_bias_lr=0.1,
+            # Loss weights
+            box=7.5,
+            cls=0.5,
+            dfl=1.5,
+            # Augmentation settings for X-ray images
+            hsv_h=0.0,  # No hue augmentation for X-ray
+            hsv_s=0.0,  # No saturation augmentation
+            hsv_v=0.1,  # Slight value augmentation
+            degrees=0.0,  # No rotation
+            translate=0.1,
+            scale=0.5,
+            shear=0.0,
+            perspective=0.0,
+            flipud=0.0,  # No vertical flip for X-ray
+            fliplr=0.5,  # Horizontal flip is okay
+            mosaic=1.0,
+            mixup=0.0,
+            copy_paste=0.0
         )
         
         # Collect training result plots
@@ -304,13 +339,22 @@ def train_model(epochs, batch_size, img_size, device_selection):
         model_path = os.path.join(results_path, "weights", "best.pt")
         
         # Load the trained model to ensure it's ready for inference
+        model_loaded = False
+        class_info = ""
         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)} 클래스)"
+            try:
+                model = YOLO(model_path)
+                model_loaded = True
+                class_info = f"\n✅ Trained on {len(model.names)} classes: {', '.join(list(model.names.values()))}"
+                
+                # Run a test inference to ensure model works
+                test_img = np.zeros((640, 640, 3), dtype=np.uint8)
+                test_results = model(test_img, verbose=False)
+                class_info += "\n✅ Model test passed - ready for inference!"
+            except Exception as e:
+                class_info = f"\n⚠️ Model loaded but test failed: {str(e)}"
         else:
-            class_info = ""
+            class_info = "\n❌ Model file not found!"
         
         training_in_progress = False
         
@@ -368,51 +412,101 @@ def run_inference(input_image, conf_threshold):
     try:
         # Check if model is trained on X-ray dataset
         model_info = ""
+        num_classes = 80  # Default COCO classes
+        class_list = []
         try:
-            model_info = f"Using model with {len(model.names)} classes\n"
+            num_classes = len(model.names)
+            class_list = list(model.names.values())
+            model_info = f"Model: {num_classes} classes - {', '.join(class_list)}\n"
+            model_info += f"Confidence threshold: {conf_threshold}\n\n"
         except:
             model_info = "Using loaded model\n"
         
-        # Save the input image temporarily
+        # Save the input image temporarily with proper format
         temp_path = "temp_inference.jpg"
-        input_image.save(temp_path)
-        
-        # Run inference with workers=0
-        results = model(temp_path, conf=conf_threshold, verbose=False, device=0 if torch.cuda.is_available() else 'cpu')
+        # Ensure image is in RGB format
+        if input_image.mode != 'RGB':
+            input_image = input_image.convert('RGB')
+        input_image.save(temp_path, format='JPEG', quality=95)
+        
+        # Print image info for debugging
+        print(f"Image mode: {input_image.mode}, Size: {input_image.size}")
+        
+        # Run inference with verbose output for debugging
+        imgsz = 640  # Processing size
+        print(f"Running inference with conf={conf_threshold}, imgsz={imgsz}")
+        results = model(
+            temp_path, 
+            conf=conf_threshold, 
+            verbose=True,  # Enable verbose for debugging
+            device=0 if torch.cuda.is_available() else 'cpu',
+            imgsz=imgsz,  # Ensure consistent image size
+            augment=False,  # Disable augmentation for consistent results
+            agnostic_nms=False,  # Use class-specific NMS
+            max_det=300  # Maximum detections
+        )
         
-        # Draw results
-        annotated_image = results[0].plot()
+        # Draw results with lower confidence for visualization
+        annotated_image = results[0].plot(
+            conf=True,  # Show confidence scores
+            labels=True,  # Show labels
+            boxes=True,  # Show boxes
+            masks=False,  # No masks for detection
+            probs=False  # No classification probabilities
+        )
         
         # Get detection info
         detections = []
-        if results[0].boxes is not None and len(results[0].boxes) > 0:
-            for box in results[0].boxes:
+        detection_count = 0
+        
+        if results[0].boxes is not None:
+            detection_count = len(results[0].boxes)
+            print(f"Found {detection_count} detections")
+            
+            for idx, box in enumerate(results[0].boxes):
                 cls = int(box.cls)
                 conf = float(box.conf)
+                xyxy = box.xyxy[0].tolist()
                 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
-            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."
+                
+                detections.append(f"{idx+1}. {cls_name}: {conf:.3f} | Box: [{int(xyxy[0])}, {int(xyxy[1])}, {int(xyxy[2])}, {int(xyxy[3])}]")
         
         # Clean up
         if os.path.exists(temp_path):
             os.remove(temp_path)
         
+        # Prepare detection text
+        if detections:
+            detection_text = model_info + f"✅ Found {detection_count} object(s):\n\n" + "\n".join(detections)
+        else:
+            # Provide helpful debugging info
+            detection_text = model_info + "❌ No objects detected.\n\n"
+            
+            if num_classes == 80:  # COCO model
+                detection_text += "⚠️ This is a COCO model (general objects). For X-ray detection, please train on the X-ray dataset.\n"
+            else:
+                detection_text += "💡 Suggestions:\n"
+                detection_text += f"• Current threshold: {conf_threshold} - try lowering to 0.1-0.3\n"
+                detection_text += "• Ensure you're using X-ray images (grayscale/blue-tinted)\n"
+                detection_text += "• Check if the model was trained with the same image size\n"
+                detection_text += "• The model may need more training epochs\n"
+                
+                # Show what the model is looking for
+                detection_text += f"\n\n📊 Debug Info:\n"
+                detection_text += f"• Image size: {input_image.size}\n"
+                detection_text += f"• Processing size: {imgsz}x{imgsz}\n"
+                detection_text += f"• Device: {'GPU' if torch.cuda.is_available() else 'CPU'}\n"
+                detection_text += f"• Model classes: {num_classes} - {', '.join(class_list)}"
+        
         return Image.fromarray(annotated_image), detection_text
         
     except Exception as e:
+        print(f"Error during inference: {str(e)}")
+        import traceback
+        traceback.print_exc()
         return None, f"Error during inference: {str(e)}"
 
 @spaces.GPU(duration=60)  # Request GPU for batch inference
@@ -454,19 +548,24 @@ def batch_inference(data_type, num_images):
             return [], f"No images found in {image_dir}"
         
         results_images = []
+        detection_counts = []
         
         for img_path in image_files:
-            results = model(img_path, verbose=False)
+            results = model(img_path, verbose=False, conf=0.25, imgsz=640)
             annotated = results[0].plot()
             results_images.append(Image.fromarray(annotated))
+            
+            # Count detections
+            if results[0].boxes is not None:
+                detection_counts.append(len(results[0].boxes))
+            else:
+                detection_counts.append(0)
         
         # Check model type
-        try:
-            model_type = "X-ray detection model" if len(model.names) != 80 else "General COCO model"
-        except:
-            model_type = "Loaded model"
+        model_type = "X-ray detection model" if len(model.names) != 80 else "General COCO model"
+        avg_detections = sum(detection_counts) / len(detection_counts) if detection_counts else 0
         
-        return results_images, f"Processed {len(results_images)} images from {data_type} dataset using {model_type}"
+        return results_images, f"Processed {len(results_images)} images using {model_type}\nAverage detections per image: {avg_detections:.1f}"
         
     except Exception as e:
         return [], f"Error during batch inference: {str(e)}"
@@ -487,10 +586,19 @@ def get_dataset_info():
         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."
+        # Count images
+        train_images = len(glob(f"{dataset_path}/train/images/*")) if os.path.exists(f"{dataset_path}/train/images") else 0
+        valid_images = len(glob(f"{dataset_path}/valid/images/*")) if os.path.exists(f"{dataset_path}/valid/images") else 0
+        test_images = len(glob(f"{dataset_path}/test/images/*")) if os.path.exists(f"{dataset_path}/test/images") else 0
+        
+        info = f"### 📊 X-ray Baggage Dataset Info\n\n"
+        info += f"**Classes ({num_classes}):** {', '.join(class_names)}\n\n"
+        info += f"**Dataset Split:**\n"
+        info += f"- Training: {train_images} images\n"
+        info += f"- Validation: {valid_images} images\n"
+        info += f"- Test: {test_images} images\n"
+        info += f"- Total: {train_images + valid_images + test_images} images\n\n"
+        info += f"**What to expect:** The model will learn to detect these prohibited items in X-ray scans."
         
         return info
     except Exception as e:
@@ -577,20 +685,21 @@ def check_model_status():
         if os.path.exists(trained_path):
             try:
                 model = YOLO(trained_path)
-                return f"✅ Trained model loaded: {len(model.names)} classes"
+                num_classes = len(model.names)
+                class_names = ', '.join(list(model.names.values()))
+                return f"✅ Trained model loaded: {num_classes} classes\n📋 Classes: {class_names}"
             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)
+            class_names = ', '.join(list(model.names.values()))
+            
             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}"
+                return f"✅ Model loaded: {num_classes} classes\n📋 Classes: {class_names}"
         except:
             return "✅ Model loaded"
 
@@ -706,18 +815,19 @@ with gr.Blocks(title="X-ray Baggage Anomaly Detection", theme=gr.themes.Soft())
         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
+        **Recommended Settings for X-ray Detection:**
+        - **Epochs:** 20-30 for good results
+        - **Batch Size:** 2-4 for better convergence
+        - **Image Size:** 640 for best quality
+        - **Expected time:** ~2-5 minutes for 20 epochs
         
         ⚠️ **Important**: Models are temporary on Spaces! Download your model after training.
         """)
         
         with gr.Row():
-            epochs_input = gr.Slider(1, 50, 10, step=1, label="Epochs")
-            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")
+            epochs_input = gr.Slider(1, 50, 20, step=1, label="Epochs (20+ recommended)")
+            batch_size_input = gr.Slider(2, 16, 4, step=2, label="Batch Size (lower for better results)")
+            img_size_input = gr.Slider(320, 640, 640, step=32, label="Image Size (640 recommended)")
             device_input = gr.Radio(["Auto", "GPU", "CPU"], value="Auto", label="Device")
         
         train_btn = gr.Button("Start Training", variant="primary")
@@ -775,40 +885,72 @@ with gr.Blocks(title="X-ray Baggage Anomaly Detection", theme=gr.themes.Soft())
         
         refresh_status_btn.click(check_model_status, outputs=model_status)
         
+        gr.Markdown("""
+        ## 🎯 모델이 객체를 감지하지 못하나요?
+        
+        **권장 학습 설정:**
+        - **Epochs: 30** (최소 20 이상)
+        - **Batch Size: 2 또는 4**
+        - **Image Size: 640**
+        
+        **체크리스트:**
+        1. ✅ X-ray 이미지인가? (일반 사진은 작동 안 함)
+        2. ✅ 충분히 학습했나? (20+ epochs)
+        3. ✅ Confidence threshold를 0.01로 낮춰봤나?
+        4. ✅ 모델이 제대로 로드되었나? (상태 확인)
+        
+        **성공적인 학습 후 예상 결과:**
+        - Firearm (총기류) 감지
+        - Knife (칼) 감지
+        - Pliers (펜치) 감지
+        - Scissors (가위) 감지
+        - Wrench (렌치) 감지
+        """)
+        
         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 X-ray Image")
-                conf_threshold = gr.Slider(0.1, 0.9, 0.5, step=0.05, label="Confidence Threshold")
+                conf_threshold = gr.Slider(0.01, 0.9, 0.25, step=0.01, label="Confidence Threshold (낮을수록 더 많이 감지)")
+                
+                # Debug options
+                with gr.Row():
+                    inference_btn = gr.Button("Run Detection", variant="primary")
+                    test_btn = gr.Button("Test with 0.01 threshold", variant="secondary", scale=0)
                 
                 # Add example images if dataset is available
-                def get_example_images():
-                    examples = []
-                    if dataset_path and os.path.exists(f"{dataset_path}/test/images"):
-                        test_images = glob(f"{dataset_path}/test/images/*")[:3]
-                        examples.extend(test_images)
-                    return examples if examples else None
+                example_images = []
+                if dataset_path and os.path.exists(f"{dataset_path}/test/images"):
+                    test_images = glob(f"{dataset_path}/test/images/*")[:5]
+                    example_images.extend(test_images)
                 
-                example_images = get_example_images()
                 if example_images:
                     gr.Examples(
-                        examples=example_images,
+                        examples=[[img] for img in example_images],
                         inputs=input_image,
-                        label="Example X-ray Images"
+                        label="Example X-ray Images (Click to load)"
                     )
-                
-                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)
+                detection_info = gr.Textbox(label="Detection Info", lines=8)
         
         inference_btn.click(run_inference, 
                           inputs=[input_image, conf_threshold],
                           outputs=[output_image, detection_info])
         
+        # Test with very low threshold
+        test_btn.click(
+            lambda img: run_inference(img, 0.01),
+            inputs=[input_image],
+            outputs=[output_image, detection_info]
+        )
+        
+        # Auto-refresh model status after inference
+        inference_btn.click(check_model_status, outputs=model_status)
+        
         gr.Markdown("### Batch Inference")
         gr.Markdown("Run detection on multiple images from the test dataset.")
         
@@ -828,7 +970,8 @@ with gr.Blocks(title="X-ray Baggage Anomaly Detection", theme=gr.themes.Soft())
     gr.Markdown("---")
     gr.Markdown("""
     <div style='text-align: center; font-size: 14px; color: #666;'>
-        💡 <b>Quick Start:</b> Download Dataset → Train Model (or Upload existing) → Run Inference<br>
+        💡 <b>Quick Start:</b> Download Dataset → Train Model (20+ epochs) → Run Inference<br>
+        🔍 <b>No detections?</b> Try lowering threshold to 0.01 or train for more epochs<br>
         🚀 Built with Gradio, YOLOv8, and ❤️ for X-ray security
     </div>
     """)