Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,83 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import streamlit as st
|
2 |
+
import numpy as np
|
3 |
+
import cv2
|
4 |
+
import joblib
|
5 |
+
import tensorflow as tf
|
6 |
+
from tensorflow.keras.applications import (
|
7 |
+
ResNet50, VGG16, EfficientNetV2B0, InceptionV3, ResNet101, DenseNet201
|
8 |
+
)
|
9 |
+
from tensorflow.keras.preprocessing import image
|
10 |
+
import os
|
11 |
+
|
12 |
+
# Define available models
|
13 |
+
MODELS = {
|
14 |
+
'ResNet50': ResNet50(weights='imagenet', include_top=False, input_shape=(224, 224, 3), pooling="avg"),
|
15 |
+
'VGG16': VGG16(weights='imagenet', include_top=False, input_shape=(224, 224, 3), pooling="avg"),
|
16 |
+
'EfficientNetV2B0': EfficientNetV2B0(weights='imagenet', include_top=False, input_shape=(224, 224, 3), pooling="avg"),
|
17 |
+
'InceptionV3': InceptionV3(weights='imagenet', include_top=False, input_shape=(224, 224, 3), pooling="avg"),
|
18 |
+
'ResNet101': ResNet101(weights='imagenet', include_top=False, input_shape=(224, 224, 3), pooling="avg"),
|
19 |
+
'DenseNet201': DenseNet201(weights='imagenet', include_top=False, input_shape=(224, 224, 3), pooling="avg")
|
20 |
+
}
|
21 |
+
|
22 |
+
# Load trained models
|
23 |
+
MODEL_PATHS = {model_name: f"{model_name}_catboost.pkl" for model_name in MODELS}
|
24 |
+
trained_models = {}
|
25 |
+
|
26 |
+
# Load the trained models into memory
|
27 |
+
for model_name, path in MODEL_PATHS.items():
|
28 |
+
if os.path.exists(path):
|
29 |
+
trained_models[model_name] = joblib.load(path)
|
30 |
+
|
31 |
+
# Define class names (modify based on dataset)
|
32 |
+
CLASS_NAMES = train.class_names # Update with actual labels
|
33 |
+
|
34 |
+
# Streamlit UI
|
35 |
+
st.title("Multi-Model Image Classifier")
|
36 |
+
st.markdown("Upload an image and select which models and classes to use for prediction.")
|
37 |
+
|
38 |
+
# Upload an image
|
39 |
+
uploaded_file = st.file_uploader("Upload an image...", type=["jpg", "png", "jpeg"])
|
40 |
+
|
41 |
+
# Select Models
|
42 |
+
selected_models = st.multiselect("Select models for prediction:", list(trained_models.keys()))
|
43 |
+
|
44 |
+
# Select Classes
|
45 |
+
selected_classes = st.multiselect("Select classes to predict:", CLASS_NAMES, default=CLASS_NAMES)
|
46 |
+
|
47 |
+
# Function to preprocess image
|
48 |
+
def preprocess_image(img_path):
|
49 |
+
img = image.load_img(img_path, target_size=(224, 224))
|
50 |
+
img_array = image.img_to_array(img) / 255.0 # Normalize
|
51 |
+
img_array = np.expand_dims(img_array, axis=0)
|
52 |
+
return img_array
|
53 |
+
|
54 |
+
# Perform prediction
|
55 |
+
if uploaded_file and selected_models:
|
56 |
+
# Read and preprocess the image
|
57 |
+
image_np = np.array(bytearray(uploaded_file.read()), dtype=np.uint8)
|
58 |
+
image_np = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
|
59 |
+
image_np = cv2.resize(image_np, (224, 224)) / 255.0
|
60 |
+
image_np = np.expand_dims(image_np, axis=0)
|
61 |
+
|
62 |
+
# Extract Features
|
63 |
+
extracted_features = {}
|
64 |
+
for model_name in selected_models:
|
65 |
+
extracted_features[model_name] = MODELS[model_name].predict(image_np)
|
66 |
+
|
67 |
+
# Predict using each selected CatBoost model
|
68 |
+
predictions = {}
|
69 |
+
for model_name in selected_models:
|
70 |
+
X_input = extracted_features[model_name]
|
71 |
+
catboost_model = trained_models[model_name]
|
72 |
+
y_pred = catboost_model.predict_proba(X_input)[0] # Get probabilities
|
73 |
+
predictions[model_name] = y_pred
|
74 |
+
|
75 |
+
# Display results
|
76 |
+
st.subheader("Prediction Results")
|
77 |
+
|
78 |
+
for model_name, y_pred in predictions.items():
|
79 |
+
st.write(f"**Model: {model_name}**")
|
80 |
+
for i, class_name in enumerate(CLASS_NAMES):
|
81 |
+
if class_name in selected_classes:
|
82 |
+
st.write(f" - {class_name}: **{y_pred[i]:.4f}**")
|
83 |
+
|