import streamlit as st import requests import io import base64 from PIL import Image, ImageOps import numpy as np import json import time import pandas as pd from typing import List from datetime import datetime # =================== CONFIG =================== API_URL = "https://b3rian-image-classifier-api.hf.space/predict" SUPPORTED_FORMATS = ["jpg", "jpeg", "png", "webp"] MAX_SIZE_MB = 10 MAX_SIZE_BYTES = MAX_SIZE_MB * 1024 * 1024 # =================== UTILITY FUNCTIONS =================== def compress_image(image: Image.Image, quality: int = 85) -> bytes: with io.BytesIO() as output: image.save(output, format='JPEG', quality=quality) return output.getvalue() def create_thumbnail(image: Image.Image, size=(128, 128)) -> str: image.thumbnail(size) with io.BytesIO() as buffer: image.save(buffer, format="JPEG", quality=70) return base64.b64encode(buffer.getvalue()).decode() def validate_image(file) -> Image.Image: try: if hasattr(file, 'size') and file.size > MAX_SIZE_BYTES: st.error(f"File too large (max {MAX_SIZE_MB}MB)") return None image = Image.open(file) image.verify() image = Image.open(file) return image.convert("RGB") except Exception as e: st.error(f"Invalid image: {str(e)}") return None def fetch_image_from_url(url: str) -> Image.Image: try: with st.spinner("Fetching image from URL..."): head_response = requests.head(url, timeout=20, allow_redirects=True) if head_response.status_code != 200: raise ValueError(f"URL returned {head_response.status_code}") response = requests.get(url, timeout=10) response.raise_for_status() return Image.open(io.BytesIO(response.content)).convert("RGB") except Exception as e: st.error(f"URL Error: {str(e)}") return None def get_image_metadata(img: Image.Image) -> str: return f"Size: {img.size}, Mode: {img.mode}, Format: {img.format}" def classify_image_with_retry(image: Image.Image, model_name: str, max_retries=2): img_bytes = compress_image(image) files = {"file": ("image.jpg", img_bytes, "image/jpeg")} params = {"model_name": model_name} for attempt in range(max_retries + 1): try: with st.spinner(f"Classifying with {model_name}..."): res = requests.post(API_URL, files=files, params=params, timeout=120) res.raise_for_status() return res.json() except requests.exceptions.ConnectionError: if attempt == max_retries: st.error("⚠️ The model server is currently offline. Please try again later.") return None time.sleep(1) except requests.exceptions.Timeout: if attempt == max_retries: st.error("⏳ The request to the model server timed out. Please try again.") return None time.sleep(1) except requests.exceptions.HTTPError as e: st.error(f"🚫 HTTP error: {e.response.status_code} - {e.response.reason}") return None except requests.exceptions.RequestException: if attempt == max_retries: st.error("🚨 An unexpected error occurred while contacting the model server.") return None time.sleep(1) def display_predictions(predictions, model_version, inference_time): st.subheader(f"Predictions: {model_version}") if not predictions: st.warning("No predictions above the confidence threshold.") return df = pd.DataFrame(predictions) df = df.set_index("label") for pred in predictions: st.markdown(f"**{pred['label']}**: {pred['confidence']}%") st.progress(pred['confidence'] / 100.0) st.caption(f"Inference time: {inference_time:.2f}s") # =================== MAIN APP =================== def main(): st.markdown("---") st.set_page_config(page_title="Image Classifier", layout="wide", page_icon="🖼️") st.title("🖼️ AI Image Classifier") st.caption("Powered by Convolutional Neural Networks (CNNs)") st.markdown(""" 📌 Upload or capture an image and choose a CNN model to classify it. 🔍 **How it works**: The selected AI model analyzes your image and returns its best predictions, sorted by confidence. """) # Initialize session state st.session_state.setdefault("history", []) st.session_state.setdefault("feedback", {}) st.session_state.setdefault("model_cache", {}) # Sidebar controls with st.sidebar: st.markdown("---") st.markdown("### ⚙️ Preferences & Model Selection") with st.expander("Advanced Options"): num_predictions = st.slider( "Number of predictions", 1, 10, 3, help="""Set how many predictions to display (1-10). Higher values show more alternatives but may include less relevant results.""" ) confidence_threshold = st.slider( "Confidence threshold (%)", 0, 100, 0, help="""Minimum confidence percentage (0-100%) required to show a prediction. Increase to filter out low-confidence results.""" ) compare_models = st.checkbox( "🔁 Compare Models", help="Run both models on the image and compare their predictions." ) model_name = st.selectbox( "Select 🧠 AI Model", ["efficientnet", "resnet"], disabled=compare_models, help="""Choose a deep learning architecture: • **EfficientNet:** Lightweight and fast (good for mobile/edge devices) • **ResNet:** Powerful general-purpose model (best accuracy/speed balance). Disabled when 'Compare Models' is active - all models will run simultaneously.""" ) st.markdown("---") st.subheader("💬 Feedback") with st.form("feedback_form_sidebar"): history = st.session_state["history"] if history: selected = st.selectbox("Select image to review", [h["name"] for h in history], help="""Choose a previously classified image to provide feedback on. The model's predictions for this image will be shown below for reference. Only images with existing classification results appear here.""") rating = st.select_slider("Rating (1-5)", options=[1, 2, 3, 4, 5], value=3, help="""Rate the model's accuracy for this image: 1 = Completely wrong • 2 = Mostly incorrect • 3 = Partially correct 4 = Mostly accurate • 5 = Perfect prediction """) selected_item = next((h for h in history if h["name"] == selected), None) if selected_item: st.markdown("**Model Predictions:**") for pred in selected_item["predictions"]: st.markdown(f"- {pred['label']}: {pred['confidence']:.1f}%") correction = st.text_input("Suggested correction", placeholder="Correct label", help="""If the AI's prediction was wrong, please provide: • The accurate label for this image • Be specific (e.g., 'Golden Retriever' instead of just 'Dog') • Use singular nouns where applicable Your input helps train better models!""") comment = st.text_area("Additional comments", placeholder="Anything else?", help="""Share details to improve the model: • What features did the AI miss? • Was the mistake understandable? • Any edge cases we should know about? (Examples: 'The turtle was partially obscured' or 'Confused labrador with golden retriever')""") else: st.info("No images classified yet.") selected = rating = correction = comment = None if st.form_submit_button("Submit Feedback", type='primary') and selected: st.session_state["feedback"][selected] = { "rating": rating, "predictions": selected_item.get("predictions", []), "correction": correction, "comment": comment, "timestamp": time.strftime("%Y-%m-%d %H:%M:%S") } st.toast("Feedback saved!", icon="✅") # Image input methods images = [] tab1, tab2, tab3 = st.tabs(["📤 Upload Image", "📷 Use Webcam", "🌐 From URL"]) with tab1: uploaded_files = st.file_uploader("Upload Image(s)", type=SUPPORTED_FORMATS, accept_multiple_files=True) for file in uploaded_files: img = validate_image(file) if img: images.append((img, file.name)) with tab2: try: picture = st.camera_input("Capture Image") if picture: img = validate_image(picture) if img: images.append((img, f"webcam_{time.strftime('%Y%m%d_%H%M%S')}.jpg")) except Exception: st.error("Webcam not supported on this device.") with tab3: url = st.text_input("Image URL", placeholder="https://example.com/image.jpg") col1, col2 = st.columns([3, 1]) if col1.button("Fetch Image", type='primary') and url: img = fetch_image_from_url(url) if img: images.append((img, f"url_{time.strftime('%Y%m%d_%H%M%S')}.jpg")) if col2.button("Clear URL", type='primary'): url = "" # Classify images if images: st.subheader("🖼️ Image Preview") for idx, (img, name) in enumerate(images): with st.expander(f"Image: {name}", expanded=True): col1, col2 = st.columns([1, 2]) with col1: st.image(img, caption=name, use_container_width=True) with col2: st.markdown(get_image_metadata(img)) if st.button("🚀 Classify Image", key=f"classify_{idx}", type='primary'): models_to_run = ["efficientnet", "resnet"] if compare_models else [model_name] for model in models_to_run: cache_key = f"{name}_{model}" result = st.session_state.model_cache.get(cache_key) if result: st.toast(f"Using cached result for {model}") else: result = classify_image_with_retry(img, model) if result: st.session_state.model_cache[cache_key] = result if result: preds = [p for p in result['predictions'] if p['confidence'] >= confidence_threshold][:num_predictions] display_predictions(preds, result['model_version'], result['inference_time']) st.session_state.history.append({ "name": name, "predictions": preds, "model": result['model_version'], "time": result.get('timestamp', datetime.now().isoformat()), "thumbnail": create_thumbnail(img) }) # Show history st.divider() st.subheader("📜 Session History") if not st.session_state.history: st.info("No classification history.") else: for record in reversed(st.session_state.history[-5:]): with st.container(border=True): col1, col2 = st.columns([1, 4]) with col1: if "thumbnail" in record: st.image(io.BytesIO(base64.b64decode(record["thumbnail"]))) with col2: st.markdown(f"**{record['name']}**") st.caption(f"Model: `{record['model']}` | {record['time']}") if record['predictions']: top_pred = record['predictions'][0] st.markdown(f"**Top Prediction**: {top_pred['label']} ({top_pred['confidence']:.1f}%)") if record['name'] in st.session_state.feedback: fb = st.session_state.feedback[record['name']] st.markdown(f"Feedback: ⭐{fb['rating']}/5") if fb['correction']: st.markdown(f"*Suggested correction: {fb['correction']}*") # Download button st.download_button( "📥 Download Results as JSON", data=json.dumps(st.session_state.history, indent=2), file_name="classification_history.json", type='primary', use_container_width=True ) st.markdown("---") st.caption("Built with ❤️ using Streamlit") if __name__ == "__main__": main()