diff --git a/README.md b/README.md
index 859894e..2875539 100644
--- a/README.md
+++ b/README.md
@@ -64,22 +64,57 @@ Please look at the [Steps to Run](#steps-to-run) section for Docker instructions
```
### Run the Script inside the Container
+
+#### Option 1: Streamlit Web Interface (Recommended) 
+Run the interactive web interface:
+```sh
+streamlit run streamlit_app.py
+```
+
+The Streamlit interface provides:
+- 🖼️ **Image Upload/Selection**: Choose from sample images or upload your own
+- ⚙️ **Easy Configuration**: Select inference mode, adjust top-K predictions via UI
+- 📊 **Real-time Results**: View predictions and benchmark metrics interactively
+- 📈 **Visual Feedback**: See benchmark results in an organized table format
+
+#### Option 2: Command Line Interface
```sh
python main.py [--mode all]
```
-### Arguments
+### Arguments (CLI)
- `--image_path`: (Optional) Specifies the path to the image you want to predict.
- `--topk`: (Optional) Specifies the number of top predictions to show. Defaults to 5 if not provided.
- `--mode`: (Optional) Specifies the model's mode for exporting and running. Choices are: `onnx`, `ov`, `cpu`, `cuda`, `tensorrt`, and `all`. If not provided, it defaults to `all`.
-### Example Command
+### Example Command (CLI)
```sh
python main.py --topk 3 --mode=all --image_path="./inference/cat3.jpg"
```
This command will run predictions on the chosen image (`./inference/cat3.jpg`), show the top 3 predictions, and run all available models. Note: plot created only for `--mode=all` and results plotted and saved to `./inference/plot.png`
+## Streamlit Interface 
+The project now includes a user-friendly Streamlit web interface for running benchmarks interactively.
+
+### Interface Preview
+
+
+### Features
+- **Interactive Image Selection**: Choose from sample images or upload your own
+- **Flexible Configuration**: Select inference modes (ONNX, OpenVINO, PyTorch CPU/CUDA, TensorRT)
+- **Real-time Benchmarking**: Run benchmarks and see results instantly
+- **Visual Results**: View predictions and performance metrics in an organized format
+- **System Information**: Check available hardware (CPU/GPU) and capabilities
+
+### Benchmark Results Display
+
+
+The interface displays:
+- Top-K predictions with confidence scores
+- Benchmark metrics (average inference time, throughput)
+- Clear visual organization of results
+
## Results
### Example Input
Here is an example of the input image to run predictions and benchmarks on:
diff --git a/pyproject.toml b/pyproject.toml
index 7278992..633c757 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -10,7 +10,7 @@ readme = "README.md"
requires-python = ">=3.12"
license = {file = "LICENSE"}
authors = [
- {name = "DimaBir", email = ""}
+ {name = "DimaBir"}
]
keywords = ["pytorch", "tensorrt", "onnx", "openvino", "inference", "deep-learning"]
classifiers = [
@@ -30,9 +30,11 @@ dependencies = [
"numpy>=1.26.0",
"onnx>=1.16.0",
"onnxruntime>=1.18.0",
+ "onnxscript>=0.5.0",
"openvino>=2024.5.0",
"seaborn>=0.13.0",
"matplotlib>=3.8.0",
+ "streamlit>=1.41.0",
]
[project.optional-dependencies]
diff --git a/requirements.txt b/requirements.txt
index 81f38ad..247aff4 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -5,8 +5,10 @@ Pillow>=10.0.0
numpy>=1.26.0
onnx>=1.16.0
onnxruntime>=1.18.0
+onnxscript>=0.5.0
openvino>=2024.5.0
seaborn>=0.13.0
matplotlib>=3.8.0
pytest>=8.0.0
pytest-cov>=4.1.0
+streamlit>=1.41.0
diff --git a/streamlit_app.py b/streamlit_app.py
new file mode 100644
index 0000000..d1961df
--- /dev/null
+++ b/streamlit_app.py
@@ -0,0 +1,331 @@
+"""
+Streamlit interface for ResNet TensorRT benchmark application.
+
+This app provides a user-friendly web interface to:
+- Upload or select images for inference
+- Configure benchmark parameters
+- Run inference across different backends (PyTorch, ONNX, OpenVINO, TensorRT)
+- Display predictions and benchmark results
+"""
+
+import os
+import tempfile
+from pathlib import Path
+
+import pandas as pd
+import streamlit as st
+import torch
+from PIL import Image
+
+from common.utils import (
+ DEFAULT_IMAGE_PATH,
+ DEFAULT_ONNX_PATH,
+ DEFAULT_OV_PATH,
+ DEFAULT_TOPK,
+ INFERENCE_MODES,
+)
+from src.image_processor import ImageProcessor
+from src.model import ModelLoader
+from src.onnx_inference import ONNXInference
+from src.ov_inference import OVInference
+from src.pytorch_inference import PyTorchInference
+
+# Check CUDA availability
+CUDA_AVAILABLE = torch.cuda.is_available()
+if CUDA_AVAILABLE:
+ try:
+ import torch_tensorrt # noqa: F401
+ from src.tensorrt_inference import TensorRTInference
+ except ImportError:
+ CUDA_AVAILABLE = False
+ st.warning("torch-tensorrt not installed. TensorRT and CUDA modes will be unavailable.")
+
+
+def display_image(image_path: str):
+ """Display the input image."""
+ img = Image.open(image_path)
+ st.image(img, caption="Input Image", use_container_width=True)
+
+
+def run_inference(
+ image_path: str,
+ mode: str,
+ topk: int,
+ onnx_path: str,
+ ov_path: str,
+ debug_mode: bool = False,
+) -> dict[str, tuple[float, float]]:
+ """
+ Run inference based on selected mode.
+
+ Args:
+ image_path: Path to input image
+ mode: Inference mode (onnx, ov, cpu, cuda, tensorrt, all)
+ topk: Number of top predictions to show
+ onnx_path: Path to ONNX model
+ ov_path: Path to OpenVINO model
+ debug_mode: Enable debug logging
+
+ Returns:
+ Dictionary of benchmark results {model_name: (avg_time_ms, throughput)}
+ """
+ benchmark_results = {}
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+ # Load model and process image
+ model_loader = ModelLoader(device=device)
+ img_processor = ImageProcessor(img_path=image_path, device=device)
+ img_batch = img_processor.process_image()
+
+ # ONNX inference
+ if mode in ["onnx", "all"]:
+ with st.spinner("Running ONNX inference..."):
+ onnx_inference = ONNXInference(model_loader, onnx_path, debug_mode=debug_mode)
+ benchmark_result = onnx_inference.benchmark(img_batch)
+ predictions = onnx_inference.predict(img_batch)
+ benchmark_results["ONNX (CPU)"] = benchmark_result
+
+ if predictions is not None:
+ display_predictions(predictions, model_loader.categories, topk, "ONNX (CPU)")
+
+ # OpenVINO inference
+ if mode in ["ov", "all"]:
+ with st.spinner("Running OpenVINO inference..."):
+ ov_inference = OVInference(model_loader, ov_path, debug_mode=debug_mode)
+ benchmark_result = ov_inference.benchmark(img_batch)
+ predictions = ov_inference.predict(img_batch)
+ benchmark_results["OpenVINO (CPU)"] = benchmark_result
+
+ if predictions is not None:
+ display_predictions(predictions, model_loader.categories, topk, "OpenVINO (CPU)")
+
+ # PyTorch CPU inference
+ if mode in ["cpu", "all"]:
+ with st.spinner("Running PyTorch CPU inference..."):
+ pytorch_cpu_inference = PyTorchInference(
+ model_loader, device="cpu", debug_mode=debug_mode
+ )
+ benchmark_result = pytorch_cpu_inference.benchmark(img_batch)
+ predictions = pytorch_cpu_inference.predict(img_batch)
+ benchmark_results["PyTorch (CPU)"] = benchmark_result
+
+ if predictions is not None:
+ display_predictions(predictions, model_loader.categories, topk, "PyTorch (CPU)")
+
+ # CUDA and TensorRT inference (only if CUDA available)
+ if CUDA_AVAILABLE:
+ # PyTorch CUDA inference
+ if mode in ["cuda", "all"]:
+ with st.spinner("Running PyTorch CUDA inference..."):
+ pytorch_cuda_inference = PyTorchInference(
+ model_loader, device=device, debug_mode=debug_mode
+ )
+ benchmark_result = pytorch_cuda_inference.benchmark(img_batch)
+ predictions = pytorch_cuda_inference.predict(img_batch)
+ benchmark_results["PyTorch (CUDA)"] = benchmark_result
+
+ if predictions is not None:
+ display_predictions(
+ predictions, model_loader.categories, topk, "PyTorch (CUDA)"
+ )
+
+ # TensorRT inference
+ if mode in ["tensorrt", "all"]:
+ precisions = [torch.float16, torch.float32]
+ for precision in precisions:
+ precision_name = "FP16" if precision == torch.float16 else "FP32"
+ with st.spinner(f"Running TensorRT {precision_name} inference..."):
+ tensorrt_inference = TensorRTInference(
+ model_loader, device=device, precision=precision, debug_mode=debug_mode
+ )
+ benchmark_result = tensorrt_inference.benchmark(img_batch)
+ predictions = tensorrt_inference.predict(img_batch)
+ benchmark_results[f"TRT_{precision}"] = benchmark_result
+
+ if predictions is not None:
+ display_predictions(
+ predictions, model_loader.categories, topk, f"TensorRT {precision_name}"
+ )
+
+ return benchmark_results
+
+
+def display_predictions(prob, categories, topk: int, model_name: str):
+ """Display top-k predictions."""
+ top_indices = prob.argsort()[-topk:][::-1]
+ top_probs = prob[top_indices]
+
+ st.subheader(f"Predictions - {model_name}")
+ for i in range(topk):
+ probability = top_probs[i]
+ class_label = categories[0][int(top_indices[i])]
+ st.write(f"#{i + 1}: {int(probability * 100)}% {class_label}")
+
+
+def display_benchmark_results(results: dict[str, tuple[float, float]]):
+ """Display benchmark results in a table."""
+ st.subheader("Benchmark Results")
+
+ # Create DataFrame for display
+ data = {
+ "Model": list(results.keys()),
+ "Avg Time (ms)": [f"{results[model][0]:.2f}" for model in results.keys()],
+ "Throughput (samples/sec)": [f"{results[model][1]:.2f}" for model in results.keys()],
+ }
+ df = pd.DataFrame(data)
+
+ st.dataframe(df, use_container_width=True)
+
+ # Display metrics in columns
+ cols = st.columns(len(results))
+ for idx, (model, (avg_time, throughput)) in enumerate(results.items()):
+ with cols[idx]:
+ st.metric(label=model, value=f"{avg_time:.2f} ms", delta=f"{throughput:.1f} img/s")
+
+
+def main():
+ st.set_page_config(
+ page_title="ResNet TensorRT Benchmark",
+ page_icon="🚀",
+ layout="wide",
+ initial_sidebar_state="expanded",
+ )
+
+ st.title("🚀 ResNet TensorRT Benchmark Interface")
+ st.markdown(
+ """
+ This application provides a user-friendly interface to benchmark ResNet inference
+ across different backends: PyTorch (CPU/CUDA), ONNX, OpenVINO, and TensorRT.
+ """
+ )
+
+ # Sidebar configuration
+ st.sidebar.header("⚙️ Configuration")
+
+ # Image selection/upload
+ st.sidebar.subheader("Image Input")
+ image_source = st.sidebar.radio("Select image source:", ["Sample Images", "Upload Image"])
+
+ image_path = None
+ is_temp_file = False # Track if file should be cleaned up
+ if image_source == "Sample Images":
+ sample_images = []
+ inference_dir = Path("./inference")
+ if inference_dir.exists():
+ sample_images = [str(f) for f in inference_dir.glob("*.jpg")] + [
+ str(f) for f in inference_dir.glob("*.png")
+ ]
+
+ if sample_images:
+ selected_image = st.sidebar.selectbox("Choose a sample image:", sample_images)
+ image_path = selected_image
+ else:
+ st.sidebar.warning("No sample images found in ./inference directory")
+ image_path = DEFAULT_IMAGE_PATH
+ else:
+ uploaded_file = st.sidebar.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+ if uploaded_file is not None:
+ # Save uploaded file to temporary location
+ with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp_file:
+ tmp_file.write(uploaded_file.read())
+ image_path = tmp_file.name
+ is_temp_file = True
+
+ # Inference mode selection
+ st.sidebar.subheader("Inference Settings")
+
+ # Filter available modes based on CUDA availability
+ available_modes = INFERENCE_MODES.copy()
+ if not CUDA_AVAILABLE:
+ available_modes = [m for m in available_modes if m not in ["cuda", "tensorrt"]]
+ st.sidebar.info("CUDA/TensorRT modes unavailable (GPU not detected)")
+
+ mode = st.sidebar.selectbox(
+ "Select inference mode:",
+ available_modes,
+ index=available_modes.index("all") if "all" in available_modes else 0,
+ help="Choose which inference backend(s) to benchmark",
+ )
+
+ topk = st.sidebar.slider(
+ "Top-K predictions:",
+ min_value=1,
+ max_value=10,
+ value=DEFAULT_TOPK,
+ help="Number of top predictions to display",
+ )
+
+ # Advanced settings
+ with st.sidebar.expander("Advanced Settings"):
+ onnx_path = st.text_input("ONNX model path:", value=DEFAULT_ONNX_PATH)
+ ov_path = st.text_input("OpenVINO model path:", value=DEFAULT_OV_PATH)
+ debug_mode = st.checkbox("Enable debug mode", value=False)
+
+ # Main content
+ col1, col2 = st.columns([1, 2])
+
+ with col1:
+ st.subheader("Input Image")
+ if image_path and os.path.exists(image_path):
+ display_image(image_path)
+ else:
+ st.warning("Please select or upload an image to proceed")
+
+ with col2:
+ st.subheader("Actions")
+
+ # System info
+ with st.expander("System Information"):
+ st.write(f"**Device:** {'CUDA (GPU)' if CUDA_AVAILABLE else 'CPU'}")
+ if CUDA_AVAILABLE:
+ st.write(f"**GPU Name:** {torch.cuda.get_device_name(0)}")
+ st.write(f"**PyTorch Version:** {torch.__version__}")
+
+ # Run benchmark button
+ if st.button("▶️ Run Benchmark", type="primary", use_container_width=True):
+ if image_path and os.path.exists(image_path):
+ try:
+ st.info(f"Running benchmark with mode: **{mode}**")
+
+ # Run inference and get results
+ results = run_inference(
+ image_path=image_path,
+ mode=mode,
+ topk=topk,
+ onnx_path=onnx_path,
+ ov_path=ov_path,
+ debug_mode=debug_mode,
+ )
+
+ # Display results
+ st.success("Benchmark completed!")
+ display_benchmark_results(results)
+
+ # Clean up temporary file if uploaded
+ if is_temp_file:
+ try:
+ os.unlink(image_path)
+ except Exception:
+ pass
+
+ except Exception as e:
+ st.error(f"Error during benchmark: {str(e)}")
+ st.exception(e)
+ else:
+ st.error("Please select or upload a valid image first!")
+
+ # Footer
+ st.markdown("---")
+ st.markdown(
+ """
+
+ """,
+ unsafe_allow_html=True,
+ )
+
+
+if __name__ == "__main__":
+ main()