Spaces:
Sleeping
Sleeping
| import os | |
| import cv2 | |
| import numpy as np | |
| import pandas as pd | |
| import base64 | |
| from flask import Flask, render_template, request, jsonify | |
| from werkzeug.utils import secure_filename | |
| from io import BytesIO | |
| from PIL import Image | |
| import tensorflow as tf | |
| from tensorflow.keras.models import load_model | |
| # Suppress warnings | |
| os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' | |
| os.environ['MPLCONFIGDIR'] = '/tmp/matplotlib' | |
| app = Flask(__name__) | |
| app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024 # 16MB max | |
| app.config['UPLOAD_FOLDER'] = 'uploads' | |
| app.config['ALLOWED_EXTENSIONS'] = {'png', 'jpg', 'jpeg'} | |
| os.makedirs(app.config['UPLOAD_FOLDER'], mode=0o777, exist_ok=True) | |
| # Load model | |
| print("Loading traffic sign classification model...") | |
| model = load_model('tabela_tespit.h5', compile=False) | |
| model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy']) | |
| print("✓ Model loaded successfully!") | |
| # Load labels | |
| labels_df = pd.read_csv('labels.csv') | |
| print(f"✓ Loaded {len(labels_df)} traffic sign classes") | |
| def allowed_file(filename): | |
| return '.' in filename and filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS'] | |
| def preprocess_image(image_path): | |
| """Preprocess image for CNN model (32x32 grayscale)""" | |
| try: | |
| # Read original image for display | |
| img_original = cv2.imread(image_path) | |
| if img_original is None: | |
| raise ValueError("Could not read image") | |
| # Convert to RGB for display | |
| img_display = cv2.cvtColor(img_original, cv2.COLOR_BGR2RGB) | |
| # Read as grayscale for model | |
| img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE) | |
| print(f"Original image shape: {img.shape}") | |
| # Resize to 32x32 for model | |
| img_resized = cv2.resize(img, (32, 32)) | |
| # Apply histogram equalization | |
| img_equalized = cv2.equalizeHist(img_resized) | |
| # Normalize to [0, 1] | |
| img_normalized = img_equalized / 255.0 | |
| # Add dimensions: (1, 32, 32, 1) | |
| img_input = img_normalized.reshape(1, 32, 32, 1) | |
| print(f"Model input shape: {img_input.shape}") | |
| return img_input, img_display | |
| except Exception as e: | |
| raise ValueError(f"Failed to preprocess image: {str(e)}") | |
| def img_to_base64(img): | |
| """Convert numpy image to base64 string""" | |
| img_pil = Image.fromarray(img.astype('uint8')) | |
| buf = BytesIO() | |
| img_pil.save(buf, format='PNG') | |
| buf.seek(0) | |
| img_base64 = base64.b64encode(buf.getvalue()).decode('utf-8') | |
| return f'data:image/png;base64,{img_base64}' | |
| def index(): | |
| return render_template('index.html') | |
| def predict(): | |
| try: | |
| if 'file' not in request.files: | |
| return jsonify({'error': 'No file uploaded'}), 400 | |
| file = request.files['file'] | |
| if file.filename == '': | |
| return jsonify({'error': 'No file selected'}), 400 | |
| if not allowed_file(file.filename): | |
| return jsonify({'error': 'Invalid file type. Please upload PNG, JPG, or JPEG'}), 400 | |
| # Save file | |
| filename = secure_filename(file.filename) | |
| filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename) | |
| file.save(filepath) | |
| print(f"Processing: {filename}") | |
| # Preprocess | |
| img_input, img_display = preprocess_image(filepath) | |
| # Predict | |
| print("Making prediction...") | |
| predictions = model.predict(img_input, verbose=0) | |
| # Get top prediction | |
| class_id = np.argmax(predictions) | |
| confidence = np.max(predictions) | |
| class_name = labels_df.loc[class_id, 'Name'] | |
| # Get top 5 predictions | |
| top5_indices = np.argsort(predictions[0])[-5:][::-1] | |
| top5_predictions = [] | |
| for idx in top5_indices: | |
| top5_predictions.append({ | |
| 'class_id': int(idx), | |
| 'class_name': labels_df.loc[idx, 'Name'], | |
| 'confidence': float(predictions[0][idx]) | |
| }) | |
| # Convert image to base64 | |
| img_base64 = img_to_base64(img_display) | |
| # Clean up | |
| os.remove(filepath) | |
| result = { | |
| 'predicted_class': class_name, | |
| 'class_id': int(class_id), | |
| 'confidence': float(confidence), | |
| 'top5_predictions': top5_predictions, | |
| 'image': img_base64 | |
| } | |
| print(f"✓ Prediction: {class_name} (Confidence: {confidence:.2%})") | |
| return jsonify(result) | |
| except Exception as e: | |
| print(f"Error during prediction: {e}") | |
| import traceback | |
| traceback.print_exc() | |
| if os.path.exists(filepath): | |
| os.remove(filepath) | |
| return jsonify({'error': str(e)}), 500 | |
| def test_example(): | |
| """Test with example image""" | |
| try: | |
| example_path = 'image.jpg' | |
| if not os.path.exists(example_path): | |
| return jsonify({'error': 'Example image not found'}), 404 | |
| print(f"Testing with example: {example_path}") | |
| # Preprocess | |
| img_input, img_display = preprocess_image(example_path) | |
| # Predict | |
| print("Making prediction on example...") | |
| predictions = model.predict(img_input, verbose=0) | |
| # Get top prediction | |
| class_id = np.argmax(predictions) | |
| confidence = np.max(predictions) | |
| class_name = labels_df.loc[class_id, 'Name'] | |
| # Get top 5 predictions | |
| top5_indices = np.argsort(predictions[0])[-5:][::-1] | |
| top5_predictions = [] | |
| for idx in top5_indices: | |
| top5_predictions.append({ | |
| 'class_id': int(idx), | |
| 'class_name': labels_df.loc[idx, 'Name'], | |
| 'confidence': float(predictions[0][idx]) | |
| }) | |
| # Convert image to base64 | |
| img_base64 = img_to_base64(img_display) | |
| result = { | |
| 'predicted_class': class_name, | |
| 'class_id': int(class_id), | |
| 'confidence': float(confidence), | |
| 'top5_predictions': top5_predictions, | |
| 'image': img_base64 | |
| } | |
| print(f"✓ Example prediction: {class_name} (Confidence: {confidence:.2%})") | |
| return jsonify(result) | |
| except Exception as e: | |
| print(f"Error during example prediction: {e}") | |
| import traceback | |
| traceback.print_exc() | |
| return jsonify({'error': str(e)}), 500 | |
| if __name__ == '__main__': | |
| app.run(host='0.0.0.0', port=7860, debug=False) | |