app.py

"""CIFAR-10 Advanced CNN Image Classifier
Flask Web Application for Image Classification
"""
import os
from pathlib import Path
from werkzeug.utils import secure_filename
from flask import Flask, render_template, request, redirect, url_for, flash

# TensorFlow imports
try:
    from tensorflow.keras.models import load_model
    TENSORFLOW_AVAILABLE = True
except ImportError:
    TENSORFLOW_AVAILABLE = False

try:
    from essentials import pred_and_plot_image
    ESSENTIALS_AVAILABLE = True
except ImportError:
    ESSENTIALS_AVAILABLE = False

# Initialize Flask app
app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'static/uploads'

# Allowed file extensions
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'bmp', 'webp'}


def allowed_file(filename):
    """Check if file extension is allowed"""
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


def get_model_path():
    """Get the path to the trained model"""
    model_path = Path('model') / 'cifar10_advanced_cnn.h5'
    if not model_path.exists():
        return None
    return model_path


# Global variables
uploaded_files = []

# Define class names (from CIFAR-10)
CLASS_NAMES = [
    '0: airplane', '1: automobile', '2: bird', '3: cat', '4: deer',
    '5: dog', '6: frog', '7: horse', '8: ship', '9: truck'
]


def get_model():
    """Load the model"""
    if not TENSORFLOW_AVAILABLE:
        return None
    
    try:
        model_path = get_model_path()
        if model_path is None:
            return None
        return load_model(str(model_path), compile=False)
    except Exception as e:
        print(f"Error loading model: {str(e)}")
        return None


def cleanup_files():
    """Clean up uploaded files"""
    global uploaded_files
    for file_path in uploaded_files:
        try:
            if os.path.exists(file_path):
                os.remove(file_path)
        except:
            pass
    uploaded_files = []


@app.route('/')
def home():
    cleanup_files()
    return render_template('main.html')


@app.route('/predict', methods=['POST'])
def predict():
    global uploaded_files
    
    if not ESSENTIALS_AVAILABLE or not TENSORFLOW_AVAILABLE:
        flash('System not fully initialized', 'error')
        return redirect(url_for('home'))
    
    try:
        images = request.files.getlist("image")
        
        if not images or all(img.filename == '' for img in images):
            flash('No images selected.', 'error')
            return redirect(url_for('home'))
        
        model = get_model()
        if model is None:
            flash('Error loading model.', 'error')
            return redirect(url_for('home'))
        
        os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
        
        image_paths = []
        
        for image in images:
            if not image or image.filename == '':
                continue
            
            if not allowed_file(image.filename):
                continue
            
            try:
                filename = secure_filename(image.filename)
                import time
                timestamp = int(time.time() * 1000)
                filename = f"{timestamp}_{filename}"
                file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
                image.save(file_path)
                image_paths.append(file_path)
                uploaded_files.append(file_path)
            except:
                continue
        
        if not image_paths:
            flash('No valid images uploaded.', 'error')
            return redirect(url_for('home'))
        
        try:
            probabilities, labels, throughput, total_time = pred_and_plot_image(
                model=model,
                class_names=CLASS_NAMES,
                image_paths=image_paths
            )
        except:
            flash('Error processing images.', 'error')
            cleanup_files()
            return redirect(url_for('home'))
        
        if not probabilities:
            flash('No predictions generated.', 'error')
            cleanup_files()
            return redirect(url_for('home'))
        
        data = {
            "image_name": image_paths,
            "Probs": probabilities,
            "Labels": labels
        }
        
        latency = (1000 / throughput) if throughput > 0 else 0
        
        return render_template(
            'testingoutput.html',
            data=data,
            length=round(total_time, 4),
            tput=round(throughput, 2),
            latency=round(latency, 4)
        )
        
    except:
        flash('Error occurred.', 'error')
        cleanup_files()
        return redirect(url_for('home'))


@app.route('/goback')
def go_back():
    cleanup_files()
    return redirect(url_for('home'))


if __name__ == '__main__':
    app.run(debug=True, port=5000)