Spaces:
Sleeping
Sleeping
File size: 2,802 Bytes
a0d2afd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
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']
@app.route('/')
def index():
return render_template('index.html')
@app.route('/predict', methods=['POST'])
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)
|