Update app.py

app.py

@@ -1,722 +1,823 @@
-from flask import Flask, request, jsonify, render_template_string, send_from_directory, send_file
-from flask_cors import CORS
-import os
-import importlib.util
-import time
-import zipfile
-import shutil
-import json
-from datetime import datetime
-from predict_task1 import Predictor
-from id_mapping import mapping
-from show_stitched import *
-import cv2
-import supervision as sv
-from ultralytics import YOLO
-
-app = Flask(__name__)
-CORS(app)
-
-# Load configuration
-config_dir = os.path.abspath(os.path.dirname(__file__))
-config_path = os.path.join(config_dir, 'PC_CONFIG.py')
-spec = importlib.util.spec_from_file_location("PC_CONFIG", config_path)
-PC_CONFIG = importlib.util.module_from_spec(spec)
-spec.loader.exec_module(PC_CONFIG)
-
-HOST = PC_CONFIG.HOST
-PORT = PC_CONFIG.IMAGE_REC_PORT
-UPLOAD_FOLDER = os.path.join(PC_CONFIG.FILE_DIRECTORY, "image-rec", "images")
-DATASET_FOLDER = os.path.join(PC_CONFIG.BASE_DIR, "yolo_dataset")
-ANNOTATED_FOLDER = os.path.join(DATASET_FOLDER, "annotated_images")
-LABELS_FOLDER = os.path.join(DATASET_FOLDER, "labels")
-IMAGES_FOLDER = os.path.join(DATASET_FOLDER, "images")
-CLASS_MAPPING_FILE = os.path.join(DATASET_FOLDER, "classes.json")
-
-app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
-
-# Initialize predictor
-predictor = Predictor()
-
-# Ensure directories exist
-os.makedirs(UPLOAD_FOLDER, exist_ok=True)
-os.makedirs(DATASET_FOLDER, exist_ok=True)
-os.makedirs(ANNOTATED_FOLDER, exist_ok=True)
-os.makedirs(LABELS_FOLDER, exist_ok=True)
-os.makedirs(IMAGES_FOLDER, exist_ok=True)
-
-# Initialize class mapping file
-if not os.path.exists(CLASS_MAPPING_FILE):
-    # Create initial class mapping from id_mapping.py
-    reverse_mapping = {str(v): k for k, v in mapping.items() if v != -1 and k is not None}
-    with open(CLASS_MAPPING_FILE, 'w', encoding='utf-8') as f:
-        json.dump(reverse_mapping, f, indent=2, ensure_ascii=False)
-
-def load_class_mapping():
-    """Load class mapping from JSON file"""
-    try:
-        with open(CLASS_MAPPING_FILE, 'r', encoding='utf-8') as f:
-            return json.load(f)
-    except:
-        return {}
-
-def save_class_mapping(class_mapping):
-    """Save class mapping to JSON file"""
-    with open(CLASS_MAPPING_FILE, 'w', encoding='utf-8') as f:
-        json.dump(class_mapping, f, indent=2, ensure_ascii=False)
-
-def generate_yolo_annotation(results, detection_id, image_width, image_height, class_name):
-    """Generate YOLO format annotation string"""
-    if not results or detection_id >= len(results[0].boxes):
-        return ""
-    
-    # Get class mapping
-    class_mapping = load_class_mapping()
-    
-    # Get class ID from mapping, if not found, add it
-    class_id = None
-    for id_str, name in class_mapping.items():
-        if name == class_name:
-            class_id = int(id_str)
-            break
-    
-    if class_id is None:
-        # Add new class to mapping
-        max_id = max([int(k) for k in class_mapping.keys()]) if class_mapping else -1
-        class_id = max_id + 1
-        class_mapping[str(class_id)] = class_name
-        save_class_mapping(class_mapping)
-    
-    # Get bounding box
-    box = results[0].boxes.xyxy[detection_id]
-    x1, y1, x2, y2 = box.tolist()
-    
-    # Convert to YOLO format (normalized)
-    x_center = ((x1 + x2) / 2) / image_width
-    y_center = ((y1 + y2) / 2) / image_height
-    width = (x2 - x1) / image_width
-    height = (y2 - y1) / image_height
-    
-    confidence = results[0].boxes.conf[detection_id].item()
-    
-    # YOLO format: class_id x_center y_center width height confidence
-    return f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f} {confidence:.6f}"
-
-def save_annotated_image(image, results, detection_id, filename):
-    """Save annotated image with bounding boxes"""
-    if not results or not results[0].boxes:
-        return None
-    
-    # Create supervision annotators
-    bounding_box_annotator = sv.BoundingBoxAnnotator()
-    label_annotator = sv.LabelAnnotator()
-    
-    # Convert YOLOv8 results to supervision format
-    boxes = results[0].boxes.xyxy.cpu().numpy()
-    confidences = results[0].boxes.conf.cpu().numpy()
-    class_ids = results[0].boxes.cls.cpu().numpy().astype(int)
-    
-    # Get class names
-    class_names = [results[0].names[class_id] for class_id in class_ids]
-    
-    # Create detections
-    detections = sv.Detections(
-        xyxy=boxes,
-        confidence=confidences,
-        class_id=class_ids
-    )
-    
-    # Annotate image
-    annotated_image = bounding_box_annotator.annotate(scene=image.copy(), detections=detections)
-    annotated_image = label_annotator.annotate(
-        scene=annotated_image, 
-        detections=detections,
-        labels=[f"{class_names[i]} {confidences[i]:.2f}" for i in range(len(class_names))]
-    )
-    
-    # Save annotated image
-    annotated_path = os.path.join(ANNOTATED_FOLDER, f"annotated_{filename}")
-    cv2.imwrite(annotated_path, annotated_image)
-    
-    return annotated_path
-
-def process_file(file_path, direction, task_type):
-    """Process uploaded file and generate predictions"""
-    print("File received and saved successfully.")
-    print(f"Direction received: {direction}")
-    print(f"Task type received: {task_type}")
-
-    startTime = datetime.now()
-    
-    # Load image
-    image = cv2.imread(file_path)
-    if image is None:
-        return None
-    
-    # Perform prediction
-    class_name, results, detection_id = predictor.predict_id(file_path, task_type)
-    class_id = str(mapping.get(class_name, -1))
-    
-    if class_name and results:
-        # Generate filename
-        timestamp = int(time.time())
-        base_filename = f"{class_name}_{timestamp}"
-        
-        # Save original image to dataset
-        image_filename = f"{base_filename}.jpg"
-        dataset_image_path = os.path.join(IMAGES_FOLDER, image_filename)
-        shutil.copy2(file_path, dataset_image_path)
-        
-        # Generate and save YOLO annotation
-        h, w = image.shape[:2]
-        yolo_annotation = generate_yolo_annotation(results, detection_id, w, h, class_name)
-        if yolo_annotation:
-            txt_filename = f"{base_filename}.txt"
-            txt_path = os.path.join(LABELS_FOLDER, txt_filename)
-            with open(txt_path, 'w') as f:
-                f.write(yolo_annotation)
-        
-        # Save annotated image
-        save_annotated_image(image, results, detection_id, image_filename)
-    
-    endTime = datetime.now()
-    totalTime = (endTime - startTime).total_seconds()
-    print(f"Predicted ID: {class_id}")
-    print(f"Time taken for Predicting Image = {totalTime} s")
-    
-    return class_id
-
-# HTML template for the frontend
-HTML_TEMPLATE = """
-<!DOCTYPE html>
-<html lang="en">
-<head>
-    <meta charset="UTF-8">
-    <meta name="viewport" content="width=device-width, initial-scale=1.0">
-    <title>YOLO Image Recognition System</title>
-    <style>
-        body {
-            font-family: Arial, sans-serif;
-            max-width: 1200px;
-            margin: 0 auto;
-            padding: 20px;
-            background-color: #f5f5f5;
-        }
-        .header {
-            text-align: center;
-            margin-bottom: 30px;
-        }
-        .container {
-            display: grid;
-            grid-template-columns: 1fr 1fr;
-            gap: 20px;
-            margin-bottom: 30px;
-        }
-        .card {
-            background-color: white;
-            padding: 20px;
-            border-radius: 10px;
-            box-shadow: 0 2px 10px rgba(0,0,0,0.1);
-        }
-        .image-display {
-            text-align: center;
-        }
-        .image-display img {
-            max-width: 100%;
-            max-height: 400px;
-            border: 2px solid #ddd;
-            border-radius: 8px;
-        }
-        .status {
-            padding: 10px;
-            border-radius: 5px;
-            margin-bottom: 15px;
-            font-weight: bold;
-        }
-        .status.waiting {
-            background-color: #fff3cd;
-            color: #856404;
-        }
-        .status.updated {
-            background-color: #d4edda;
-            color: #155724;
-        }
-        .upload-section {
-            text-align: center;
-        }
-        .upload-form {
-            display: inline-block;
-            text-align: left;
-        }
-        .form-group {
-            margin-bottom: 15px;
-        }
-        .form-group label {
-            display: block;
-            margin-bottom: 5px;
-            font-weight: bold;
-        }
-        .form-group input, .form-group select {
-            width: 300px;
-            padding: 8px;
-            border: 1px solid #ddd;
-            border-radius: 4px;
-        }
-        .btn {
-            background-color: #007bff;
-            color: white;
-            border: none;
-            padding: 12px 24px;
-            border-radius: 5px;
-            font-size: 16px;
-            cursor: pointer;
-            transition: background-color 0.3s;
-        }
-        .btn:hover {
-            background-color: #0056b3;
-        }
-        .btn:disabled {
-            background-color: #6c757d;
-            cursor: not-allowed;
-        }
-        .dataset-info {
-            grid-column: span 2;
-        }
-        .stats-grid {
-            display: grid;
-            grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
-            gap: 15px;
-            margin-top: 15px;
-        }
-        .stat-card {
-            background-color: #f8f9fa;
-            padding: 15px;
-            border-radius: 5px;
-            text-align: center;
-        }
-        .stat-number {
-            font-size: 24px;
-            font-weight: bold;
-            color: #007bff;
-        }
-        .timestamp {
-            color: #666;
-            font-size: 12px;
-            margin-top: 10px;
-        }
-        .class-mapping {
-            max-height: 200px;
-            overflow-y: auto;
-            background-color: #f8f9fa;
-            padding: 10px;
-            border-radius: 5px;
-            font-family: monospace;
-            font-size: 12px;
-        }
-    </style>
-</head>
-<body>
-    <div class="header">
-        <h1>YOLO Image Recognition System</h1>
-        <p>Real-time image processing with dataset management</p>
-    </div>
-
-    <div class="container">
-        <div class="card">
-            <h3>Latest Recognition Result</h3>
-            <div id="status" class="status waiting">Waiting for recognition results...</div>
-            <div class="image-display">
-                <div id="result" style="display: none;">
-                    <img id="resultImage" />
-                    <div id="timestamp" class="timestamp"></div>
-                </div>
-                <div id="noResult" style="color: #666; padding: 40px;">
-                    No recognition results yet
-                </div>
-            </div>
-        </div>
-
-        <div class="card">
-            <h3>Upload Image for Recognition</h3>
-            <div class="upload-section">
-                <form id="uploadForm" class="upload-form" enctype="multipart/form-data">
-                    <div class="form-group">
-                        <label for="file">Select Image:</label>
-                        <input type="file" id="file" name="file" accept="image/*" required>
-                    </div>
-                    <div class="form-group">
-                        <label for="direction">Direction:</label>
-                        <select id="direction" name="direction" required>
-                            <option value="north">North</option>
-                            <option value="south">South</option>
-                            <option value="east">East</option>
-                            <option value="west">West</option>
-                        </select>
-                    </div>
-                    <div class="form-group">
-                        <label for="task_type">Task Type:</label>
-                        <select id="task_type" name="task_type" required>
-                            <option value="TASK_1">Task 1</option>
-                            <option value="TASK_2">Task 2</option>
-                        </select>
-                    </div>
-                    <button type="submit" class="btn">Upload and Predict</button>
-                </form>
-                <div id="uploadResult" style="margin-top: 15px;"></div>
-            </div>
-        </div>
-
-        <div class="card dataset-info">
-            <h3>Dataset Information</h3>
-            <div class="stats-grid">
-                <div class="stat-card">
-                    <div class="stat-number" id="totalImages">0</div>
-                    <div>Total Images</div>
-                </div>
-                <div class="stat-card">
-                    <div class="stat-number" id="totalClasses">0</div>
-                    <div>Total Classes</div>
-                </div>
-                <div class="stat-card">
-                    <div class="stat-number" id="annotatedImages">0</div>
-                    <div>Annotated Images</div>
-                </div>
-                <div class="stat-card">
-                    <button id="downloadBtn" class="btn" onclick="downloadDataset()">
-                        Download Dataset
-                    </button>
-                </div>
-            </div>
-            <div style="margin-top: 20px;">
-                <h4>Class Mapping:</h4>
-                <div id="classMapping" class="class-mapping">Loading...</div>
-            </div>
-            <div id="downloadInfo" style="margin-top: 10px; color: #666;"></div>
-        </div>
-    </div>
-
-    <script>
-        let lastImagePath = '';
-        
-        // Check for latest results
-        async function checkLatestResult() {
-            try {
-                const response = await fetch('/latest-result');
-                const data = await response.json();
-                
-                if (data.success && data.image_path) {
-                    const newImagePath = data.image_path;
-                    
-                    if (newImagePath !== lastImagePath) {
-                        lastImagePath = newImagePath;
-                        
-                        const resultDiv = document.getElementById('result');
-                        const noResultDiv = document.getElementById('noResult');
-                        const resultImg = document.getElementById('resultImage');
-                        const timestampDiv = document.getElementById('timestamp');
-                        const statusDiv = document.getElementById('status');
-                        
-                        // Update image and timestamp
-                        resultImg.src = '/annotated/' + newImagePath + '?t=' + new Date().getTime();
-                        timestampDiv.textContent = 'Last updated: ' + new Date().toLocaleString();
-                        
-                        // Show result
-                        noResultDiv.style.display = 'none';
-                        resultDiv.style.display = 'block';
-                        
-                        // Update status
-                        statusDiv.className = 'status updated';
-                        statusDiv.textContent = 'New recognition result available';
-                        
-                        setTimeout(() => {
-                            statusDiv.className = 'status waiting';
-                            statusDiv.textContent = 'Waiting for next result...';
-                        }, 3000);
-                    }
-                }
-            } catch (error) {
-                console.error('Failed to check latest result:', error);
-            }
-        }
-        
-        // Update dataset statistics
-        async function updateDatasetStats() {
-            try {
-                const response = await fetch('/dataset-stats');
-                const data = await response.json();
-                
-                document.getElementById('totalImages').textContent = data.total_images || 0;
-                document.getElementById('totalClasses').textContent = data.total_classes || 0;
-                document.getElementById('annotatedImages').textContent = data.annotated_images || 0;
-                
-                // Update class mapping
-                const mappingDiv = document.getElementById('classMapping');
-                if (data.class_mapping) {
-                    let mappingText = '';
-                    for (const [id, name] of Object.entries(data.class_mapping)) {
-                        mappingText += `${id}: ${name}\\n`;
-                    }
-                    mappingDiv.textContent = mappingText || 'No classes defined yet';
-                } else {
-                    mappingDiv.textContent = 'No classes defined yet';
-                }
-            } catch (error) {
-                console.error('Failed to update dataset stats:', error);
-            }
-        }
-        
-        // Handle file upload
-        document.getElementById('uploadForm').addEventListener('submit', async function(e) {
-            e.preventDefault();
-            
-            const formData = new FormData(this);
-            const uploadResult = document.getElementById('uploadResult');
-            const submitBtn = this.querySelector('button[type="submit"]');
-            
-            submitBtn.disabled = true;
-            submitBtn.textContent = 'Processing...';
-            uploadResult.innerHTML = '<div style="color: #007bff;">Processing image...</div>';
-            
-            try {
-                const response = await fetch('/upload', {
-                    method: 'POST',
-                    body: formData
-                });
-                
-                const result = await response.json();
-                
-                if (response.ok) {
-                    uploadResult.innerHTML = `
-                        <div style="color: #28a745;">
-                            <strong>Success!</strong><br>
-                            Predicted ID: ${result.predicted_id}
-                        </div>
-                    `;
-                    // Refresh dataset stats
-                    updateDatasetStats();
-                } else {
-                    uploadResult.innerHTML = `<div style="color: #dc3545;">Error: ${result.error}</div>`;
-                }
-            } catch (error) {
-                uploadResult.innerHTML = `<div style="color: #dc3545;">Error: ${error.message}</div>`;
-            } finally {
-                submitBtn.disabled = false;
-                submitBtn.textContent = 'Upload and Predict';
-            }
-        });
-        
-        // Download dataset
-        async function downloadDataset() {
-            const downloadBtn = document.getElementById('downloadBtn');
-            const downloadInfo = document.getElementById('downloadInfo');
-            
-            try {
-                downloadBtn.disabled = true;
-                downloadBtn.textContent = 'Preparing...';
-                downloadInfo.textContent = 'Creating ZIP file, please wait...';
-                downloadInfo.style.color = '#007bff';
-                
-                const response = await fetch('/download-dataset');
-                
-                if (response.ok) {
-                    const blob = await response.blob();
-                    const url = window.URL.createObjectURL(blob);
-                    const a = document.createElement('a');
-                    a.href = url;
-                    a.download = `yolo_dataset_${new Date().toISOString().slice(0,10)}.zip`;
-                    document.body.appendChild(a);
-                    a.click();
-                    window.URL.revokeObjectURL(url);
-                    document.body.removeChild(a);
-                    
-                    downloadInfo.textContent = 'Dataset downloaded successfully!';
-                    downloadInfo.style.color = '#28a745';
-                } else {
-                    const errorData = await response.json();
-                    downloadInfo.textContent = 'Error: ' + (errorData.message || 'Failed to download');
-                    downloadInfo.style.color = '#dc3545';
-                }
-            } catch (error) {
-                downloadInfo.textContent = 'Error: ' + error.message;
-                downloadInfo.style.color = '#dc3545';
-            } finally {
-                downloadBtn.disabled = false;
-                downloadBtn.textContent = 'Download Dataset';
-                setTimeout(() => { downloadInfo.textContent = ''; }, 5000);
-            }
-        }
-        
-        // Initialize
-        checkLatestResult();
-        updateDatasetStats();
-        
-        // Auto-refresh every 3 seconds
-        setInterval(checkLatestResult, 3000);
-        setInterval(updateDatasetStats, 10000);
-    </script>
-</body>
-</html>
-"""
-
-# Routes
-@app.route('/')
-def index():
-    """Home page with web interface"""
-    return render_template_string(HTML_TEMPLATE)
-
-@app.route('/status', methods=['GET'])
-def server_status():
-    """Health check endpoint"""
-    return jsonify({'status': 'OK'})
-
-@app.route('/upload', methods=['POST'])
-def upload_file():
-    """Handle file upload and prediction"""
-    if 'file' not in request.files:
-        return jsonify({'error': 'No file part'}), 400
-    
-    file = request.files['file']
-    direction = request.form.get('direction', 'north')
-    task_type = request.form.get('task_type', 'TASK_1')
-    
-    if file.filename == '':
-        return jsonify({'error': 'No selected file'}), 400
-    
-    if file:
-        filename = os.path.basename(file.filename)
-        file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
-        file.save(file_path)
-        
-        # Process the file and predict
-        class_id = process_file(file_path, direction, task_type)
-        
-        if class_id is not None:
-            return jsonify({
-                'message': 'File successfully uploaded and processed',
-                'predicted_id': class_id,
-                'direction': direction,
-                'task_type': task_type
-            }), 200
-        else:
-            return jsonify({'error': 'Failed to process image'}), 500
-
-@app.route('/latest-result')
-def get_latest_result():
-    """Get the latest annotated image"""
-    if not os.path.exists(ANNOTATED_FOLDER):
-        return jsonify({"success": False, "message": "Annotated folder not found"})
-    
-    # Get all annotated images
-    annotated_files = []
-    for filename in os.listdir(ANNOTATED_FOLDER):
-        if filename.startswith('annotated_') and filename.lower().endswith(('.png', '.jpg', '.jpeg')):
-            filepath = os.path.join(ANNOTATED_FOLDER, filename)
-            mtime = os.path.getmtime(filepath)
-            annotated_files.append((filename, mtime))
-    
-    if not annotated_files:
-        return jsonify({"success": False, "message": "No annotated images found"})
-    
-    # Sort by modification time, get latest
-    annotated_files.sort(key=lambda x: x[1], reverse=True)
-    latest_file = annotated_files[0][0]
-    
-    return jsonify({
-        "success": True,
-        "image_path": latest_file,
-        "timestamp": annotated_files[0][1]
-    })
-
-@app.route('/annotated/<filename>')
-def serve_annotated_image(filename):
-    """Serve annotated images"""
-    return send_from_directory(ANNOTATED_FOLDER, filename)
-
-@app.route('/dataset-stats')
-def get_dataset_stats():
-    """Get dataset statistics"""
-    stats = {
-        'total_images': 0,
-        'total_classes': 0,
-        'annotated_images': 0,
-        'class_mapping': {}
-    }
-    
-    # Count images
-    if os.path.exists(IMAGES_FOLDER):
-        stats['total_images'] = len([f for f in os.listdir(IMAGES_FOLDER) if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
-    
-    # Count annotated images
-    if os.path.exists(ANNOTATED_FOLDER):
-        stats['annotated_images'] = len([f for f in os.listdir(ANNOTATED_FOLDER) if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
-    
-    # Load class mapping
-    stats['class_mapping'] = load_class_mapping()
-    stats['total_classes'] = len(stats['class_mapping'])
-    
-    return jsonify(stats)
-
-@app.route('/download-dataset')
-def download_dataset():
-    """Download the complete YOLO dataset as ZIP"""
-    if not os.path.exists(DATASET_FOLDER):
-        return jsonify({"success": False, "message": "Dataset folder not found"}), 404
-    
-    # Check if there are files to download
-    has_files = False
-    for folder in [IMAGES_FOLDER, LABELS_FOLDER, ANNOTATED_FOLDER]:
-        if os.path.exists(folder) and os.listdir(folder):
-            has_files = True
-            break
-    
-    if not has_files:
-        return jsonify({"success": False, "message": "No files found in dataset"}), 404
-    
-    # Create timestamp for filename
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    zip_filename = f"yolo_dataset_{timestamp}.zip"
-    zip_path = os.path.join(UPLOAD_FOLDER, zip_filename)
-    
-    try:
-        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
-            # Add all files from dataset structure
-            for root, dirs, files in os.walk(DATASET_FOLDER):
-                for file in files:
-                    file_path = os.path.join(root, file)
-                    arcname = os.path.relpath(file_path, DATASET_FOLDER)
-                    zipf.write(file_path, arcname)
-        
-        return send_file(zip_path, as_attachment=True, download_name=zip_filename)
-    
-    except Exception as e:
-        return jsonify({"success": False, "message": f"Error creating zip: {str(e)}"}), 500
-    
-    finally:
-        # Clean up temporary file
-        try:
-            if os.path.exists(zip_path):
-                os.remove(zip_path)
-        except:
-            pass
-
-@app.route('/display_stitched', methods=['POST'])
-def display_stitched():
-    """Display stitched images"""
-    try:
-        showAnnotatedStitched()
-        return jsonify({'display_stitched': 'OK'})
-    except Exception as e:
-        return jsonify({'error': str(e)}), 500
-
-if __name__ == '__main__':
-    print()
-    print(f"UPLOAD FOLDER: {UPLOAD_FOLDER}")
-    print(f"DATASET FOLDER: {DATASET_FOLDER}")
-    print(f"Starting Enhanced Image Recognition Server...")
-    print(f"Web interface available at: http://{HOST}:{PORT}")
-    
-    try:
-        app.run(host=HOST, port=PORT, debug=False)
-    except:
-        print('Unable to connect to configured host and port. Switching to localhost:4000.')
-        app.run(host='0.0.0.0', port=4000, debug=True)
+from flask import Flask, request, jsonify, render_template_string, send_from_directory, send_file
+from flask_cors import CORS
+import os
+import importlib.util
+import time
+import zipfile
+import shutil
+import json
+from datetime import datetime
+from predict_task1 import Predictor
+from id_mapping import mapping
+from show_stitched import *
+import cv2
+import supervision as sv
+from ultralytics import YOLO
+
+app = Flask(__name__)
+CORS(app)
+
+# Load configuration
+config_dir = os.path.abspath(os.path.dirname(__file__))
+config_path = os.path.join(config_dir, 'PC_CONFIG.py')
+spec = importlib.util.spec_from_file_location("PC_CONFIG", config_path)
+PC_CONFIG = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(PC_CONFIG)
+
+HOST = PC_CONFIG.HOST
+PORT = 7860  # Changed from PC_CONFIG.IMAGE_REC_PORT to 7860
+UPLOAD_FOLDER = os.path.join(PC_CONFIG.FILE_DIRECTORY, "image-rec", "images")
+DATASET_FOLDER = os.path.join(PC_CONFIG.BASE_DIR, "yolo_dataset")
+ANNOTATED_FOLDER = os.path.join(DATASET_FOLDER, "annotated_images")
+LABELS_FOLDER = os.path.join(DATASET_FOLDER, "labels")
+IMAGES_FOLDER = os.path.join(DATASET_FOLDER, "images")
+CLASS_MAPPING_FILE = os.path.join(DATASET_FOLDER, "classes.json")
+
+app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
+
+# Initialize predictor
+predictor = Predictor()
+
+# Ensure directories exist
+os.makedirs(UPLOAD_FOLDER, exist_ok=True)
+os.makedirs(DATASET_FOLDER, exist_ok=True)
+os.makedirs(ANNOTATED_FOLDER, exist_ok=True)
+os.makedirs(LABELS_FOLDER, exist_ok=True)
+os.makedirs(IMAGES_FOLDER, exist_ok=True)
+
+# Initialize class mapping file
+if not os.path.exists(CLASS_MAPPING_FILE):
+    # Create initial class mapping from id_mapping.py
+    reverse_mapping = {str(v): k for k, v in mapping.items() if v != -1 and k is not None}
+    with open(CLASS_MAPPING_FILE, 'w', encoding='utf-8') as f:
+        json.dump(reverse_mapping, f, indent=2, ensure_ascii=False)
+
+def load_class_mapping():
+    """Load class mapping from JSON file"""
+    try:
+        with open(CLASS_MAPPING_FILE, 'r', encoding='utf-8') as f:
+            return json.load(f)
+    except:
+        return {}
+
+def save_class_mapping(class_mapping):
+    """Save class mapping to JSON file"""
+    with open(CLASS_MAPPING_FILE, 'w', encoding='utf-8') as f:
+        json.dump(class_mapping, f, indent=2, ensure_ascii=False)
+
+def generate_yolo_annotation(results, detection_id, image_width, image_height, class_name):
+    """Generate YOLO format annotation string"""
+    if not results or detection_id >= len(results[0].boxes):
+        return ""
+    
+    # Get class mapping
+    class_mapping = load_class_mapping()
+    
+    # Get class ID from mapping, if not found, add it
+    class_id = None
+    for id_str, name in class_mapping.items():
+        if name == class_name:
+            class_id = int(id_str)
+            break
+    
+    if class_id is None:
+        # Add new class to mapping
+        max_id = max([int(k) for k in class_mapping.keys()]) if class_mapping else -1
+        class_id = max_id + 1
+        class_mapping[str(class_id)] = class_name
+        save_class_mapping(class_mapping)
+    
+    # Get bounding box
+    box = results[0].boxes.xyxy[detection_id]
+    x1, y1, x2, y2 = box.tolist()
+    
+    # Convert to YOLO format (normalized)
+    x_center = ((x1 + x2) / 2) / image_width
+    y_center = ((y1 + y2) / 2) / image_height
+    width = (x2 - x1) / image_width
+    height = (y2 - y1) / image_height
+    
+    confidence = results[0].boxes.conf[detection_id].item()
+    
+    # YOLO format: class_id x_center y_center width height confidence
+    return f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f} {confidence:.6f}"
+
+def save_annotated_image(image, results, detection_id, filename):
+    """Save annotated image with bounding boxes"""
+    if not results or not results[0].boxes:
+        return None
+    
+    # Create supervision annotators
+    bounding_box_annotator = sv.BoundingBoxAnnotator()
+    label_annotator = sv.LabelAnnotator()
+    
+    # Convert YOLOv8 results to supervision format
+    boxes = results[0].boxes.xyxy.cpu().numpy()
+    confidences = results[0].boxes.conf.cpu().numpy()
+    class_ids = results[0].boxes.cls.cpu().numpy().astype(int)
+    
+    # Get class names
+    class_names = [results[0].names[class_id] for class_id in class_ids]
+    
+    # Create detections
+    detections = sv.Detections(
+        xyxy=boxes,
+        confidence=confidences,
+        class_id=class_ids
+    )
+    
+    # Annotate image
+    annotated_image = bounding_box_annotator.annotate(scene=image.copy(), detections=detections)
+    annotated_image = label_annotator.annotate(
+        scene=annotated_image, 
+        detections=detections,
+        labels=[f"{class_names[i]} {confidences[i]:.2f}" for i in range(len(class_names))]
+    )
+    
+    # Save annotated image
+    annotated_path = os.path.join(ANNOTATED_FOLDER, f"annotated_{filename}")
+    cv2.imwrite(annotated_path, annotated_image)
+    
+    return annotated_path
+
+def process_file(file_path, direction, task_type, filename):
+    """Process uploaded file and generate predictions"""
+    print("File received and saved successfully.")
+    print(f"Direction received: {direction}")
+    print(f"Task type received: {task_type}")
+
+    startTime = datetime.now()
+    
+    # Load image
+    image = cv2.imread(file_path)
+    if image is None:
+        return None
+    
+    # Perform prediction
+    class_name, results, detection_id = predictor.predict_id(file_path, task_type)
+    class_id = str(mapping.get(class_name, -1))
+    
+    detection_result = None
+    if class_name and results:
+        # Generate filename
+        timestamp = int(time.time())
+        base_filename = f"{class_name}_{timestamp}"
+        
+        # Save original image to dataset
+        image_filename = f"{base_filename}.jpg"
+        dataset_image_path = os.path.join(IMAGES_FOLDER, image_filename)
+        shutil.copy2(file_path, dataset_image_path)
+        
+        # Generate and save YOLO annotation
+        h, w = image.shape[:2]
+        yolo_annotation = generate_yolo_annotation(results, detection_id, w, h, class_name)
+        txt_path = None
+        if yolo_annotation:
+            txt_filename = f"{base_filename}.txt"
+            txt_path = os.path.join(LABELS_FOLDER, txt_filename)
+            with open(txt_path, 'w') as f:
+                f.write(yolo_annotation)
+        
+        # Save annotated image
+        annotated_path = save_annotated_image(image, results, detection_id, image_filename)
+        
+        # Get bounding box and confidence for compatibility
+        box = results[0].boxes.xyxy[detection_id]
+        confidence = results[0].boxes.conf[detection_id].item()
+        x1, y1, x2, y2 = box.tolist()
+        
+        # Create detection result in compatible format
+        detection_result = {
+            "image_id": class_id,
+            "label": class_name,
+            "confidence": confidence,
+            "bbox": [x1, y1, x2, y2],
+            "original_image_path": dataset_image_path,
+            "marked_image_path": annotated_path,
+            "txt_file_path": txt_path
+        }
+    
+    endTime = datetime.now()
+    totalTime = (endTime - startTime).total_seconds()
+    print(f"Predicted ID: {class_id}")
+    print(f"Time taken for Predicting Image = {totalTime} s")
+    
+    return class_id, detection_result
+
+# HTML template for the frontend
+HTML_TEMPLATE = """
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>YOLO Image Recognition System</title>
+    <style>
+        body {
+            font-family: Arial, sans-serif;
+            max-width: 1200px;
+            margin: 0 auto;
+            padding: 20px;
+            background-color: #f5f5f5;
+        }
+        .header {
+            text-align: center;
+            margin-bottom: 30px;
+        }
+        .container {
+            display: grid;
+            grid-template-columns: 1fr 1fr;
+            gap: 20px;
+            margin-bottom: 30px;
+        }
+        .card {
+            background-color: white;
+            padding: 20px;
+            border-radius: 10px;
+            box-shadow: 0 2px 10px rgba(0,0,0,0.1);
+        }
+        .image-display {
+            text-align: center;
+        }
+        .image-display img {
+            max-width: 100%;
+            max-height: 400px;
+            border: 2px solid #ddd;
+            border-radius: 8px;
+        }
+        .status {
+            padding: 10px;
+            border-radius: 5px;
+            margin-bottom: 15px;
+            font-weight: bold;
+        }
+        .status.waiting {
+            background-color: #fff3cd;
+            color: #856404;
+        }
+        .status.updated {
+            background-color: #d4edda;
+            color: #155724;
+        }
+        .upload-section {
+            text-align: center;
+        }
+        .upload-form {
+            display: inline-block;
+            text-align: left;
+        }
+        .form-group {
+            margin-bottom: 15px;
+        }
+        .form-group label {
+            display: block;
+            margin-bottom: 5px;
+            font-weight: bold;
+        }
+        .form-group input, .form-group select {
+            width: 300px;
+            padding: 8px;
+            border: 1px solid #ddd;
+            border-radius: 4px;
+        }
+        .btn {
+            background-color: #007bff;
+            color: white;
+            border: none;
+            padding: 12px 24px;
+            border-radius: 5px;
+            font-size: 16px;
+            cursor: pointer;
+            transition: background-color 0.3s;
+        }
+        .btn:hover {
+            background-color: #0056b3;
+        }
+        .btn:disabled {
+            background-color: #6c757d;
+            cursor: not-allowed;
+        }
+        .dataset-info {
+            grid-column: span 2;
+        }
+        .stats-grid {
+            display: grid;
+            grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
+            gap: 15px;
+            margin-top: 15px;
+        }
+        .stat-card {
+            background-color: #f8f9fa;
+            padding: 15px;
+            border-radius: 5px;
+            text-align: center;
+        }
+        .stat-number {
+            font-size: 24px;
+            font-weight: bold;
+            color: #007bff;
+        }
+        .timestamp {
+            color: #666;
+            font-size: 12px;
+            margin-top: 10px;
+        }
+        .class-mapping {
+            max-height: 200px;
+            overflow-y: auto;
+            background-color: #f8f9fa;
+            padding: 10px;
+            border-radius: 5px;
+            font-family: monospace;
+            font-size: 12px;
+        }
+    </style>
+</head>
+<body>
+    <div class="header">
+        <h1>YOLO Image Recognition System</h1>
+        <p>Real-time image processing with dataset management</p>
+    </div>
+
+    <div class="container">
+        <div class="card">
+            <h3>Latest Recognition Result</h3>
+            <div id="status" class="status waiting">Waiting for recognition results...</div>
+            <div class="image-display">
+                <div id="result" style="display: none;">
+                    <img id="resultImage" />
+                    <div id="timestamp" class="timestamp"></div>
+                </div>
+                <div id="noResult" style="color: #666; padding: 40px;">
+                    No recognition results yet
+                </div>
+            </div>
+        </div>
+
+        <div class="card">
+            <h3>Upload Image for Recognition</h3>
+            <div class="upload-section">
+                <form id="uploadForm" class="upload-form" enctype="multipart/form-data">
+                    <div class="form-group">
+                        <label for="file">Select Image:</label>
+                        <input type="file" id="file" name="file" accept="image/*" required>
+                    </div>
+                    <div class="form-group">
+                        <label for="direction">Direction:</label>
+                        <select id="direction" name="direction" required>
+                            <option value="north">North</option>
+                            <option value="south">South</option>
+                            <option value="east">East</option>
+                            <option value="west">West</option>
+                        </select>
+                    </div>
+                    <div class="form-group">
+                        <label for="task_type">Task Type:</label>
+                        <select id="task_type" name="task_type" required>
+                            <option value="TASK_1">Task 1</option>
+                            <option value="TASK_2">Task 2</option>
+                        </select>
+                    </div>
+                    <button type="submit" class="btn">Upload and Predict</button>
+                </form>
+                <div id="uploadResult" style="margin-top: 15px;"></div>
+            </div>
+        </div>
+
+        <div class="card dataset-info">
+            <h3>Dataset Information</h3>
+            <div class="stats-grid">
+                <div class="stat-card">
+                    <div class="stat-number" id="totalImages">0</div>
+                    <div>Total Images</div>
+                </div>
+                <div class="stat-card">
+                    <div class="stat-number" id="totalClasses">0</div>
+                    <div>Total Classes</div>
+                </div>
+                <div class="stat-card">
+                    <div class="stat-number" id="annotatedImages">0</div>
+                    <div>Annotated Images</div>
+                </div>
+                <div class="stat-card">
+                    <button id="downloadBtn" class="btn" onclick="downloadDataset()">
+                        Download Dataset
+                    </button>
+                </div>
+            </div>
+            <div style="margin-top: 20px;">
+                <h4>Class Mapping:</h4>
+                <div id="classMapping" class="class-mapping">Loading...</div>
+            </div>
+            <div id="downloadInfo" style="margin-top: 10px; color: #666;"></div>
+        </div>
+    </div>
+
+    <script>
+        let lastImagePath = '';
+        
+        // Check for latest results
+        async function checkLatestResult() {
+            try {
+                const response = await fetch('/latest-result');
+                const data = await response.json();
+                
+                if (data.success && data.image_path) {
+                    const newImagePath = data.image_path;
+                    
+                    if (newImagePath !== lastImagePath) {
+                        lastImagePath = newImagePath;
+                        
+                        const resultDiv = document.getElementById('result');
+                        const noResultDiv = document.getElementById('noResult');
+                        const resultImg = document.getElementById('resultImage');
+                        const timestampDiv = document.getElementById('timestamp');
+                        const statusDiv = document.getElementById('status');
+                        
+                        // Update image and timestamp
+                        resultImg.src = '/annotated/' + newImagePath + '?t=' + new Date().getTime();
+                        timestampDiv.textContent = 'Last updated: ' + new Date().toLocaleString();
+                        
+                        // Show result
+                        noResultDiv.style.display = 'none';
+                        resultDiv.style.display = 'block';
+                        
+                        // Update status
+                        statusDiv.className = 'status updated';
+                        statusDiv.textContent = 'New recognition result available';
+                        
+                        setTimeout(() => {
+                            statusDiv.className = 'status waiting';
+                            statusDiv.textContent = 'Waiting for next result...';
+                        }, 3000);
+                    }
+                }
+            } catch (error) {
+                console.error('Failed to check latest result:', error);
+            }
+        }
+        
+        // Update dataset statistics
+        async function updateDatasetStats() {
+            try {
+                const response = await fetch('/dataset-stats');
+                const data = await response.json();
+                
+                document.getElementById('totalImages').textContent = data.total_images || 0;
+                document.getElementById('totalClasses').textContent = data.total_classes || 0;
+                document.getElementById('annotatedImages').textContent = data.annotated_images || 0;
+                
+                // Update class mapping
+                const mappingDiv = document.getElementById('classMapping');
+                if (data.class_mapping) {
+                    let mappingText = '';
+                    for (const [id, name] of Object.entries(data.class_mapping)) {
+                        mappingText += `${id}: ${name}\\n`;
+                    }
+                    mappingDiv.textContent = mappingText || 'No classes defined yet';
+                } else {
+                    mappingDiv.textContent = 'No classes defined yet';
+                }
+            } catch (error) {
+                console.error('Failed to update dataset stats:', error);
+            }
+        }
+        
+        // Handle file upload
+        document.getElementById('uploadForm').addEventListener('submit', async function(e) {
+            e.preventDefault();
+            
+            const formData = new FormData(this);
+            const uploadResult = document.getElementById('uploadResult');
+            const submitBtn = this.querySelector('button[type="submit"]');
+            
+            submitBtn.disabled = true;
+            submitBtn.textContent = 'Processing...';
+            uploadResult.innerHTML = '<div style="color: #007bff;">Processing image...</div>';
+            
+            try {
+                const response = await fetch('/image', {
+                    method: 'POST',
+                    body: formData
+                });
+                
+                const result = await response.json();
+                
+                if (response.ok) {
+                    // Handle both legacy format (predicted_id) and new format (result.image_id)
+                    const predictedId = result.predicted_id || result.image_id || result.result?.image_id;
+                    const obstacleId = result.obstacle_id || result.result?.obstacle_id;
+                    
+                    uploadResult.innerHTML = `
+                        <div style="color: #28a745;">
+                            <strong>Success!</strong><br>
+                            ${obstacleId ? `Obstacle ID: ${obstacleId}<br>` : ''}
+                            Predicted ID: ${predictedId}
+                        </div>
+                    `;
+                    // Refresh dataset stats
+                    updateDatasetStats();
+                } else {
+                    uploadResult.innerHTML = `<div style="color: #dc3545;">Error: ${result.error}</div>`;
+                }
+            } catch (error) {
+                uploadResult.innerHTML = `<div style="color: #dc3545;">Error: ${error.message}</div>`;
+            } finally {
+                submitBtn.disabled = false;
+                submitBtn.textContent = 'Upload and Predict';
+            }
+        });
+        
+        // Download dataset
+        async function downloadDataset() {
+            const downloadBtn = document.getElementById('downloadBtn');
+            const downloadInfo = document.getElementById('downloadInfo');
+            
+            try {
+                downloadBtn.disabled = true;
+                downloadBtn.textContent = 'Preparing...';
+                downloadInfo.textContent = 'Creating ZIP file, please wait...';
+                downloadInfo.style.color = '#007bff';
+                
+                const response = await fetch('/download-dataset');
+                
+                if (response.ok) {
+                    const blob = await response.blob();
+                    const url = window.URL.createObjectURL(blob);
+                    const a = document.createElement('a');
+                    a.href = url;
+                    a.download = `yolo_dataset_${new Date().toISOString().slice(0,10)}.zip`;
+                    document.body.appendChild(a);
+                    a.click();
+                    window.URL.revokeObjectURL(url);
+                    document.body.removeChild(a);
+                    
+                    downloadInfo.textContent = 'Dataset downloaded successfully!';
+                    downloadInfo.style.color = '#28a745';
+                } else {
+                    const errorData = await response.json();
+                    downloadInfo.textContent = 'Error: ' + (errorData.message || 'Failed to download');
+                    downloadInfo.style.color = '#dc3545';
+                }
+            } catch (error) {
+                downloadInfo.textContent = 'Error: ' + error.message;
+                downloadInfo.style.color = '#dc3545';
+            } finally {
+                downloadBtn.disabled = false;
+                downloadBtn.textContent = 'Download Dataset';
+                setTimeout(() => { downloadInfo.textContent = ''; }, 5000);
+            }
+        }
+        
+        // Initialize
+        checkLatestResult();
+        updateDatasetStats();
+        
+        // Auto-refresh every 3 seconds
+        setInterval(checkLatestResult, 3000);
+        setInterval(updateDatasetStats, 10000);
+    </script>
+</body>
+</html>
+"""
+
+# Routes
+@app.route('/')
+def index():
+    """Home page with web interface"""
+    return render_template_string(HTML_TEMPLATE)
+
+@app.route('/status', methods=['GET'])
+def server_status():
+    """Health check endpoint"""
+    return jsonify({'status': 'OK'})
+
+@app.route('/upload', methods=['POST'])
+def upload_file():
+    """Handle file upload and prediction (legacy endpoint)"""
+    if 'file' not in request.files:
+        return jsonify({'error': 'No file part'}), 400
+    
+    file = request.files['file']
+    direction = request.form.get('direction', 'north')
+    task_type = request.form.get('task_type', 'TASK_1')
+    
+    if file.filename == '':
+        return jsonify({'error': 'No selected file'}), 400
+    
+    if file:
+        filename = os.path.basename(file.filename)
+        file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
+        file.save(file_path)
+        
+        # Process the file and predict
+        class_id, detection_result = process_file(file_path, direction, task_type, filename)
+        
+        if class_id is not None:
+            return jsonify({
+                'message': 'File successfully uploaded and processed',
+                'predicted_id': class_id,
+                'direction': direction,
+                'task_type': task_type
+            }), 200
+        else:
+            return jsonify({'error': 'Failed to process image'}), 500
+
+@app.route('/image', methods=['POST'])
+def image_predict():
+    """
+    This is the main endpoint for the image prediction algorithm
+    :return: a json object with a key "result" and value a dictionary with keys "obstacle_id" and "image_id"
+    """
+    if 'file' not in request.files:
+        return jsonify({'error': 'No file part'}), 400
+    
+    file = request.files['file']
+    filename = file.filename
+    
+    if filename == '':
+        return jsonify({'error': 'No selected file'}), 400
+    
+    # Save to uploads folder first
+    file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
+    file.save(file_path)
+    
+    # Try to parse filename format: "<timestamp>_<obstacle_id>_<signal>.jpeg"
+    # But be flexible with different formats
+    constituents = file.filename.split("_")
+    
+    # Default values
+    obstacle_id = "unknown"
+    signal = "C"  # Default to center
+    
+    # Try to extract obstacle_id and signal if available
+    try:
+        if len(constituents) >= 2:
+            obstacle_id = constituents[1]
+        if len(constituents) >= 3:
+            # Remove file extension from signal
+            signal_part = constituents[2]
+            # Handle both .jpg and .png extensions
+            for ext in ['.jpg', '.jpeg', '.png', '.JPG', '.JPEG', '.PNG']:
+                if signal_part.endswith(ext):
+                    signal = signal_part[:-len(ext)]
+                    break
+            else:
+                signal = signal_part
+    except IndexError:
+        # Use default values if parsing fails
+        pass
+    
+    # Check for optional preference parameter
+    prefer_close = request.form.get('prefer_close_objects', 'true').lower() == 'true'
+    task_type = request.form.get('task_type', 'TASK_1')
+    
+    # Process the file and predict
+    class_id, detection_result = process_file(file_path, signal, task_type, filename)
+    
+    if detection_result is None:
+        return jsonify({'error': 'Failed to process image'}), 500
+    
+    # Extract image_id from detection result
+    image_id = detection_result["image_id"]
+    
+    print(f"Original image saved to: {detection_result.get('original_image_path', 'N/A')}")
+    print(f"Annotated image saved to: {detection_result['marked_image_path']}")
+    print(f"YOLO txt file saved to: {detection_result.get('txt_file_path', 'N/A')}")
+    
+    # Return detailed detection information in compatible format
+    result = {
+        "obstacle_id": obstacle_id,
+        "image_id": image_id,
+        "detection": {
+            "label": detection_result["label"],
+            "confidence": detection_result["confidence"],
+            "bbox_coordinates": detection_result["bbox"],
+            "original_image_path": detection_result.get("original_image_path"),
+            "annotated_image_path": detection_result["marked_image_path"],
+            "txt_file_path": detection_result.get("txt_file_path")
+        }
+    }
+    return jsonify(result)
+
+@app.route('/latest-result')
+def get_latest_result():
+    """Get the latest annotated image"""
+    if not os.path.exists(ANNOTATED_FOLDER):
+        return jsonify({"success": False, "message": "Annotated folder not found"})
+    
+    # Get all annotated images
+    annotated_files = []
+    for filename in os.listdir(ANNOTATED_FOLDER):
+        if filename.startswith('annotated_') and filename.lower().endswith(('.png', '.jpg', '.jpeg')):
+            filepath = os.path.join(ANNOTATED_FOLDER, filename)
+            mtime = os.path.getmtime(filepath)
+            annotated_files.append((filename, mtime))
+    
+    if not annotated_files:
+        return jsonify({"success": False, "message": "No annotated images found"})
+    
+    # Sort by modification time, get latest
+    annotated_files.sort(key=lambda x: x[1], reverse=True)
+    latest_file = annotated_files[0][0]
+    
+    return jsonify({
+        "success": True,
+        "image_path": latest_file,
+        "timestamp": annotated_files[0][1]
+    })
+
+@app.route('/annotated/<filename>')
+def serve_annotated_image(filename):
+    """Serve annotated images"""
+    return send_from_directory(ANNOTATED_FOLDER, filename)
+
+@app.route('/dataset-stats')
+def get_dataset_stats():
+    """Get dataset statistics"""
+    stats = {
+        'total_images': 0,
+        'total_classes': 0,
+        'annotated_images': 0,
+        'class_mapping': {}
+    }
+    
+    # Count images
+    if os.path.exists(IMAGES_FOLDER):
+        stats['total_images'] = len([f for f in os.listdir(IMAGES_FOLDER) if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
+    
+    # Count annotated images
+    if os.path.exists(ANNOTATED_FOLDER):
+        stats['annotated_images'] = len([f for f in os.listdir(ANNOTATED_FOLDER) if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
+    
+    # Load class mapping
+    stats['class_mapping'] = load_class_mapping()
+    stats['total_classes'] = len(stats['class_mapping'])
+    
+    return jsonify(stats)
+
+@app.route('/download-dataset')
+def download_dataset():
+    """Download the complete YOLO dataset as ZIP"""
+    if not os.path.exists(DATASET_FOLDER):
+        return jsonify({"success": False, "message": "Dataset folder not found"}), 404
+    
+    # Check if there are files to download
+    has_files = False
+    for folder in [IMAGES_FOLDER, LABELS_FOLDER, ANNOTATED_FOLDER]:
+        if os.path.exists(folder) and os.listdir(folder):
+            has_files = True
+            break
+    
+    if not has_files:
+        return jsonify({"success": False, "message": "No files found in dataset"}), 404
+    
+    # Create timestamp for filename
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    zip_filename = f"yolo_dataset_{timestamp}.zip"
+    zip_path = os.path.join(UPLOAD_FOLDER, zip_filename)
+    
+    try:
+        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+            # Add all files from dataset structure
+            for root, dirs, files in os.walk(DATASET_FOLDER):
+                for file in files:
+                    file_path = os.path.join(root, file)
+                    arcname = os.path.relpath(file_path, DATASET_FOLDER)
+                    zipf.write(file_path, arcname)
+        
+        return send_file(zip_path, as_attachment=True, download_name=zip_filename)
+    
+    except Exception as e:
+        return jsonify({"success": False, "message": f"Error creating zip: {str(e)}"}), 500
+    
+    finally:
+        # Clean up temporary file
+        try:
+            if os.path.exists(zip_path):
+                os.remove(zip_path)
+        except:
+            pass
+
+@app.route('/display_stitched', methods=['POST'])
+def display_stitched():
+    """Display stitched images"""
+    try:
+        showAnnotatedStitched()
+        return jsonify({'display_stitched': 'OK'})
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+if __name__ == '__main__':
+    print()
+    print(f"UPLOAD FOLDER: {UPLOAD_FOLDER}")
+    print(f"DATASET FOLDER: {DATASET_FOLDER}")
+    print(f"Starting Enhanced Image Recognition Server...")
+    print(f"Web interface available at: http://{HOST}:{PORT}")
+    print(f"API endpoints: /image (main), /upload (legacy)")
+    
+    try:
+        app.run(host=HOST, port=PORT, debug=False)
+    except:
+        print('Unable to connect to configured host and port. Switching to localhost:7860.')
+        app.run(host='0.0.0.0', port=7860, debug=True)