Spaces:
Sleeping
Sleeping
| import os | |
| import secrets | |
| from flask import Flask, render_template, request, jsonify | |
| from PIL import Image | |
| import numpy as np | |
| # Try to import ML libraries, but don't crash if they are missing | |
| HAS_ML = False | |
| try: | |
| import tensorflow as tf | |
| from tensorflow.keras.models import load_model | |
| # Check if model exists | |
| model_path = os.path.join(os.path.dirname(__file__), 'cifar10_cnn_v1.h5') | |
| if os.path.exists(model_path): | |
| model = load_model(model_path) | |
| HAS_ML = True | |
| print("ML Model loaded successfully.") | |
| except Exception as e: | |
| print(f"ML mode disabled: {e}") | |
| app = Flask(__name__) | |
| app.config['UPLOAD_FOLDER'] = 'uploads' | |
| app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024 # 16MB | |
| # CIFAR-10 classes | |
| CLASSES = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'] | |
| def index(): | |
| return render_template('index.html') | |
| def predict(): | |
| if 'file' not in request.files: | |
| return jsonify({'success': False, 'error': 'No file part'}) | |
| file = request.files['file'] | |
| if file.filename == '': | |
| return jsonify({'success': False, 'error': 'No selected file'}) | |
| try: | |
| # Save file | |
| filename = secrets.token_hex(8) + "_" + file.filename | |
| filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename) | |
| file.save(filepath) | |
| # Process image | |
| img = Image.open(filepath).convert('RGB') | |
| img_resized = img.resize((32, 32)) | |
| if HAS_ML: | |
| # Real Inference | |
| img_array = np.array(img_resized) / 255.0 | |
| img_array = np.expand_dims(img_array, axis=0) | |
| predictions = model.predict(img_array) | |
| class_idx = np.argmax(predictions[0]) | |
| confidence = float(predictions[0][class_idx]) | |
| class_name = CLASSES[class_idx] | |
| else: | |
| # Mock Inference for demonstration if environment is broken | |
| # We use the filename hash to pick a "random" but consistent class for the same image | |
| hash_val = sum(ord(c) for c in filename) | |
| class_idx = hash_val % len(CLASSES) | |
| class_name = CLASSES[class_idx] | |
| confidence = 0.85 + (hash_val % 15) / 100.0 | |
| return jsonify({ | |
| 'success': True, | |
| 'class': class_name, | |
| 'confidence': confidence, | |
| 'mode': 'real' if HAS_ML else 'mock' | |
| }) | |
| except Exception as e: | |
| return jsonify({'success': False, 'error': str(e)}) | |
| if __name__ == '__main__': | |
| if not os.path.exists(app.config['UPLOAD_FOLDER']): | |
| os.makedirs(app.config['UPLOAD_FOLDER']) | |
| app.run(debug=True, port=5000) | |