diff --git a/main.py b/main.py index b878140..797b7ab 100644 --- a/main.py +++ b/main.py @@ -37,7 +37,7 @@ MODEL_PATH = 'runs/detect/memory_module_detection/weights/best.pt' detector = MemoryModuleDetector(MODEL_PATH) # Hardcoded test image path -HARDCODED_IMAGE_PATH = 'training/memory/out1.png' +HARDCODED_IMAGE_PATH = 'training/no_memory/out1.png' def allowed_file(filename): """Check if file extension is allowed.""" @@ -199,10 +199,10 @@ def detect_memory_modules(): def detect_hardcoded_image(): """ Process hardcoded test image for memory module detection. - + Optional query parameters: - - confidence: confidence threshold (default: 0.5) - + - confidence: confidence threshold (default: 0.8) + Returns: - JSON with detections and annotated image (base64) """ @@ -213,26 +213,26 @@ def detect_hardcoded_image(): 'error': 'Model not loaded. Please train the model first.', 'success': False }), 500 - + # Check if hardcoded image exists if not os.path.exists(HARDCODED_IMAGE_PATH): return jsonify({ 'error': f'Hardcoded test image not found at {HARDCODED_IMAGE_PATH}', 'success': False }), 404 - + # Get confidence threshold from query parameters (default 80%) conf_threshold = float(request.args.get('confidence', 0.8)) - + # Run detection detections, annotated_image = detector.detect( - HARDCODED_IMAGE_PATH, + HARDCODED_IMAGE_PATH, conf_threshold=conf_threshold ) - + # Convert annotated image to base64 annotated_base64 = image_to_base64(annotated_image) - + # Prepare response response_data = { 'success': True, @@ -242,9 +242,9 @@ def detect_hardcoded_image(): 'confidence_threshold': conf_threshold, 'test_image_path': HARDCODED_IMAGE_PATH } - + logger.info(f"Processed hardcoded image: found {len(detections)} memory modules") - + return jsonify(response_data) except Exception as e: