Update app.py

added Progress Bar and ETA

app.py

@@ -1,284 +1,305 @@
-import streamlit as st
-import tempfile
-import pandas as pd
-import numpy as np
-from PIL import Image
-from ultralytics import YOLO
-
-def save_uploaded_file(uploaded_file):
-    """Save an uploaded file to a temporary file and return its path."""
-    with tempfile.NamedTemporaryFile(delete=False, suffix=uploaded_file.name) as tmp_file:
-        tmp_file.write(uploaded_file.getbuffer())
-        return tmp_file.name
-
-def yolo_inference_tool():
-    """
-    Single-model, single-image inference subpage (example).
-    """
-    st.header("YOLO Model Inference Tool")
-    st.write("Upload an image and a YOLO model (.pt) file to run inference and view detailed results.")
-
-    image_file = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"], key="inference_image")
-    model_file = st.file_uploader("Upload YOLO model (.pt)", type=["pt"], key="inference_model")
-
-    if st.button("Submit (Single-Model Inference)"):
-        if not image_file or not model_file:
-            st.error("Please upload both an image and a model.")
-            return
-
-        # Save files
-        image_path = save_uploaded_file(image_file)
-        model_path = save_uploaded_file(model_file)
-
-        # Load image
-        try:
-            image = Image.open(image_file).convert("RGB")
-        except Exception as e:
-            st.error(f"Error reading image: {e}")
-            return
-
-        st.subheader("Image Details")
-        st.write(f"**Image Size:** {image.size[0]} x {image.size[1]}")
-        st.write(f"**File Type:** {image_file.type}")
-
-        # Load model
-        try:
-            model = YOLO(model_path)
-        except Exception as e:
-            st.error(f"Error loading model: {e}")
-            return
-
-        # Inference
-        st.subheader("Inference Results")
-        try:
-            results = model(np.array(image))
-        except Exception as e:
-            st.error(f"Inference error: {e}")
-            return
-
-        r = results[0]
-        boxes_data = []
-        if r.boxes is not None and len(r.boxes) > 0:
-            for i in range(len(r.boxes)):
-                coords = r.boxes.xyxy[i].cpu().numpy() if hasattr(r.boxes.xyxy[i], "cpu") else r.boxes.xyxy[i]
-                conf = r.boxes.conf[i].cpu().numpy() if hasattr(r.boxes.conf[i], "cpu") else r.boxes.conf[i]
-                cls_idx = int(r.boxes.cls[i].cpu().numpy()) if hasattr(r.boxes.cls[i], "cpu") else int(r.boxes.cls[i])
-                class_name = r.names.get(cls_idx, "Unknown")
-                boxes_data.append({
-                    "Box": i + 1,
-                    "Coordinates": f"[{coords[0]:.1f}, {coords[1]:.1f}, {coords[2]:.1f}, {coords[3]:.1f}]",
-                    "Confidence": f"{conf:.2f}",
-                    "Class": class_name
-                })
-            df_boxes = pd.DataFrame(boxes_data)
-            st.subheader("Detected Objects")
-            st.dataframe(df_boxes, use_container_width=True)
-        else:
-            st.write("No objects detected.")
-
-        # Annotated image
-        try:
-            annotated_img_bgr = r.plot(conf=True, boxes=True, labels=True)
-            annotated_img_rgb = Image.fromarray(annotated_img_bgr[..., ::-1])
-            st.subheader("Annotated Image")
-            st.image(annotated_img_rgb, caption="Inference Output", use_container_width=True)
-        except Exception as e:
-            st.error(f"Error generating annotated image: {e}")
-
-def yolo_model_comparison_tool():
-    """
-    Multi-model, multi-image comparison subpage,
-    with Weighted Scoring that uses a reciprocal-based speed metric.
-    """
-    st.header("YOLO Models Comparison Tool (Multi-Image, Weighted Score)")
-    st.write(
-        "Upload **one or more images** and **multiple YOLO model (.pt) files**. "
-        "Then click **Submit** to run inference across all images with each model. "
-        "We aggregate metrics (Avg Inference Time, Total Detections, Avg Confidence) "
-        "and compute a Weighted Score that balances these factors.\n\n"
-        "**Speed** is handled by converting inference time to a 'speed' value: `1/time`. "
-        "That way, the fastest model is near 1, and slower models get fractions in (0,1), "
-        "rather than forcing the slowest model to 0."
-    )
-
-    images = st.file_uploader("Upload Images", type=["jpg", "jpeg", "png"], key="comparison_images", accept_multiple_files=True)
-    model_files = st.file_uploader("Upload YOLO models (.pt)", type=["pt"], key="comparison_models", accept_multiple_files=True)
-
-    # Example weights. You can expose them as sliders if you want user customization.
-    alpha_detection = 0.4
-    beta_confidence = 0.3
-    gamma_speed = 0.3  # speed = reciprocal of time
-
-    if st.button("Submit (Multi-Model Comparison)"):
-        if not images or not model_files:
-            st.error("Please upload at least one image and at least one model.")
-            return
-
-        # We'll store aggregated metrics here
-        model_agg_data = {}
-        # We'll store results for each (model, image) so we can display side-by-side
-        model_image_results = {m.name: {} for m in model_files}
-
-        for model_file in model_files:
-            model_path = save_uploaded_file(model_file)
-            try:
-                model = YOLO(model_path)
-            except Exception as e:
-                st.error(f"Error loading model {model_file.name}: {e}")
-                continue
-
-            total_inference_time = 0.0
-            total_detections = 0
-            sum_confidences = 0.0
-            total_conf_count = 0
-
-            for img_file in images:
-                try:
-                    pil_img = Image.open(img_file).convert("RGB")
-                    np_img = np.array(pil_img)
-                except Exception as e:
-                    st.error(f"Error reading image {img_file.name}: {e}")
-                    continue
-
-                # Run inference
-                try:
-                    result = model(np_img)
-                except Exception as e:
-                    st.error(f"Inference error for model {model_file.name} on {img_file.name}: {e}")
-                    continue
-
-                r = result[0]
-                model_image_results[model_file.name][img_file.name] = r
-
-                # Accumulate inference time
-                if isinstance(r.speed, dict) and "inference" in r.speed:
-                    total_inference_time += r.speed["inference"]
-
-                # Count detections & confidence
-                if r.boxes is not None:
-                    det_count = len(r.boxes)
-                    total_detections += det_count
-                    if det_count > 0:
-                        confs = r.boxes.conf.cpu().numpy() if hasattr(r.boxes.conf, "cpu") else r.boxes.conf
-                        sum_confidences += confs.sum()
-                        total_conf_count += det_count
-
-            # After all images for this model
-            image_count = len(images)
-            avg_inference_time = total_inference_time / image_count if image_count > 0 else float("inf")
-            avg_confidence = sum_confidences / total_conf_count if total_conf_count > 0 else 0.0
-
-            model_agg_data[model_file.name] = {
-                "Model File": model_file.name,
-                "Avg Inference Time (ms)": avg_inference_time,
-                "Total Detections": total_detections,
-                "Average Confidence": avg_confidence
-            }
-
-        if not model_agg_data:
-            st.write("No valid models processed.")
-            return
-
-        # Display aggregated metrics
-        df = pd.DataFrame(model_agg_data.values())
-        st.subheader("Aggregated Metrics (Across All Images)")
-        st.dataframe(df, use_container_width=True)
-
-        # Weighted Scoring with reciprocal-based speed
-        # 1) Normalize detection & confidence as usual
-        detection_max = df["Total Detections"].max()
-        confidence_max = df["Average Confidence"].max()
-        if detection_max == 0: detection_max = 1
-        if confidence_max == 0: confidence_max = 1
-
-        df["Detection Norm"] = df["Total Detections"] / detection_max
-        df["Confidence Norm"] = df["Average Confidence"] / confidence_max
-
-        # 2) Convert time to speed = 1 / time, then normalize
-        #    If time = 0, add small epsilon to avoid dividing by zero
-        eps = 1e-9
-        df["Speed Val"] = 1.0 / (df["Avg Inference Time (ms)"] + eps)
-        max_speed_val = df["Speed Val"].max() if not df["Speed Val"].isnull().all() else 1
-        if max_speed_val == 0:
-            max_speed_val = 1  # fallback
-
-        df["Speed Norm"] = df["Speed Val"] / max_speed_val
-
-        # 3) Weighted Score
-        df["Weighted Score"] = (
-            alpha_detection * df["Detection Norm"] +
-            beta_confidence * df["Confidence Norm"] +
-            gamma_speed     * df["Speed Norm"]
-        )
-
-        st.subheader("Weighted Score Analysis")
-        st.write(f"Weights: Detection={alpha_detection}, Confidence={beta_confidence}, Speed={gamma_speed}")
-        st.dataframe(df[[
-            "Model File",
-            "Avg Inference Time (ms)",
-            "Total Detections",
-            "Average Confidence",
-            "Detection Norm",
-            "Confidence Norm",
-            "Speed Val",
-            "Speed Norm",
-            "Weighted Score"
-        ]], use_container_width=True)
-
-        # Identify best overall model (highest Weighted Score)
-        best_idx = df["Weighted Score"].idxmax()
-        best_model = df.loc[best_idx, "Model File"]
-        best_score = df.loc[best_idx, "Weighted Score"]
-
-        st.markdown(f"""
-        **Best Overall Model** based on Weighted Score:
-        **{best_model}** (Score: {best_score:.3f}).
-
-        ### Interpretation:
-        - **Detection Norm** → fraction of the best detection count.
-        - **Confidence Norm** → fraction of the highest average confidence.
-        - **Speed Norm** → fraction of the highest (1/time). The fastest model is near 1; others are a fraction of that speed.
-
-        If you find one factor more important, adjust the weights:
-        - Increase **Detection** weight if you care about finding as many objects as possible.
-        - Increase **Confidence** weight if you only trust high‐confidence detections.
-        - Increase **Speed** weight if you need real‐time inference.
-        """)
-
-        # Display annotated images in a grid (row = image, column = model)
-        st.subheader("Annotated Images Grid (Row = Image, Column = Model)")
-        model_names_sorted = sorted(model_agg_data.keys())
-
-        for img_file in images:
-            st.markdown(f"### Image: {img_file.name}")
-            columns = st.columns(len(model_names_sorted))
-            for col, model_name in zip(columns, model_names_sorted):
-                # Retrieve the result object
-                r = model_image_results.get(model_name, {}).get(img_file.name, None)
-                if r is None:
-                    col.write(f"No results for {model_name}")
-                    continue
-
-                # Generate annotated image
-                try:
-                    annotated_img_bgr = r.plot(conf=True, boxes=True, labels=True)
-                    annotated_img_rgb = Image.fromarray(annotated_img_bgr[..., ::-1])
-                    col.image(
-                        annotated_img_rgb,
-                        caption=f"{model_name}",
-                        use_container_width=True
-                    )
-                except Exception as e:
-                    col.error(f"Error annotating image for {model_name}: {e}")
-
-def main():
-    st.sidebar.title("Navigation")
-    page = st.sidebar.radio("Go to", ("YOLO Model Inference Tool", "YOLO Models Comparison Tool"))
-
-    if page == "YOLO Model Inference Tool":
-        yolo_inference_tool()
-    else:
-        yolo_model_comparison_tool()
-
-if __name__ == "__main__":
-    main()
+import streamlit as st
+import tempfile
+import pandas as pd
+import numpy as np
+import time
+from PIL import Image
+from ultralytics import YOLO
+
+def save_uploaded_file(uploaded_file):
+    """Save an uploaded file to a temporary file and return its path."""
+    with tempfile.NamedTemporaryFile(delete=False, suffix=uploaded_file.name) as tmp_file:
+        tmp_file.write(uploaded_file.getbuffer())
+        return tmp_file.name
+
+def yolo_inference_tool():
+    """
+    Single-model, single-image inference subpage (example).
+    """
+    st.header("YOLO Model Inference Tool")
+    st.write("Upload an image and a YOLO model (.pt) file to run inference and view detailed results.")
+
+    image_file = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"], key="inference_image")
+    model_file = st.file_uploader("Upload YOLO model (.pt)", type=["pt"], key="inference_model")
+
+    if st.button("Submit (Single-Model Inference)"):
+        if not image_file or not model_file:
+            st.error("Please upload both an image and a model.")
+            return
+
+        # Save files
+        image_path = save_uploaded_file(image_file)
+        model_path = save_uploaded_file(model_file)
+
+        # Load image
+        try:
+            image = Image.open(image_file).convert("RGB")
+        except Exception as e:
+            st.error(f"Error reading image: {e}")
+            return
+
+        st.subheader("Image Details")
+        st.write(f"**Image Size:** {image.size[0]} x {image.size[1]}")
+        st.write(f"**File Type:** {image_file.type}")
+
+        # Load model
+        try:
+            model = YOLO(model_path)
+        except Exception as e:
+            st.error(f"Error loading model: {e}")
+            return
+
+        # Inference
+        st.subheader("Inference Results")
+        try:
+            results = model(np.array(image))
+        except Exception as e:
+            st.error(f"Inference error: {e}")
+            return
+
+        r = results[0]
+        boxes_data = []
+        if r.boxes is not None and len(r.boxes) > 0:
+            for i in range(len(r.boxes)):
+                coords = r.boxes.xyxy[i].cpu().numpy() if hasattr(r.boxes.xyxy[i], "cpu") else r.boxes.xyxy[i]
+                conf = r.boxes.conf[i].cpu().numpy() if hasattr(r.boxes.conf[i], "cpu") else r.boxes.conf[i]
+                cls_idx = int(r.boxes.cls[i].cpu().numpy()) if hasattr(r.boxes.cls[i], "cpu") else int(r.boxes.cls[i])
+                class_name = r.names.get(cls_idx, "Unknown")
+                boxes_data.append({
+                    "Box": i + 1,
+                    "Coordinates": f"[{coords[0]:.1f}, {coords[1]:.1f}, {coords[2]:.1f}, {coords[3]:.1f}]",
+                    "Confidence": f"{conf:.2f}",
+                    "Class": class_name
+                })
+            df_boxes = pd.DataFrame(boxes_data)
+            st.subheader("Detected Objects")
+            st.dataframe(df_boxes, use_container_width=True)
+        else:
+            st.write("No objects detected.")
+
+        # Annotated image
+        try:
+            annotated_img_bgr = r.plot(conf=True, boxes=True, labels=True)
+            annotated_img_rgb = Image.fromarray(annotated_img_bgr[..., ::-1])
+            st.subheader("Annotated Image")
+            st.image(annotated_img_rgb, caption="Inference Output", use_container_width=True)
+        except Exception as e:
+            st.error(f"Error generating annotated image: {e}")
+
+def yolo_model_comparison_tool():
+    """
+    Multi-model, multi-image comparison subpage,
+    with Weighted Scoring that uses a reciprocal-based speed metric
+    and a real-time progress bar + ETA display.
+    """
+    st.header("YOLO Models Comparison Tool (Multi-Image, Weighted Score + Progress Bar)")
+    st.write(
+        "Upload **one or more images** and **multiple YOLO model (.pt) files**. "
+        "Then click **Submit** to run inference across all images with each model. "
+        "We aggregate metrics (Avg Inference Time, Total Detections, Avg Confidence) "
+        "and compute a Weighted Score that balances these factors.\n\n"
+        "A progress bar and ETA are shown in real time after you click Submit."
+    )
+
+    images = st.file_uploader("Upload Images", type=["jpg", "jpeg", "png"], key="comparison_images", accept_multiple_files=True)
+    model_files = st.file_uploader("Upload YOLO models (.pt)", type=["pt"], key="comparison_models", accept_multiple_files=True)
+
+    # Example weights. You can expose them as sliders if you want user customization.
+    alpha_detection = 0.4
+    beta_confidence = 0.3
+    gamma_speed = 0.3  # speed = reciprocal of time
+
+    if st.button("Submit (Multi-Model Comparison)"):
+        if not images or not model_files:
+            st.error("Please upload at least one image and at least one model.")
+            return
+
+        # Initialize progress tracking
+        total_inferences = len(images) * len(model_files)
+        if total_inferences == 0:
+            st.error("No valid images or models to process.")
+            return
+
+        progress_bar = st.progress(0)
+        eta_placeholder = st.empty()
+        start_time = time.time()
+        steps_done = 0
+
+        # We'll store aggregated metrics here
+        model_agg_data = {}
+        # We'll store results for each (model, image) so we can display side-by-side
+        model_image_results = {m.name: {} for m in model_files}
+
+        for model_file in model_files:
+            model_path = save_uploaded_file(model_file)
+            try:
+                model = YOLO(model_path)
+            except Exception as e:
+                st.error(f"Error loading model {model_file.name}: {e}")
+                continue
+
+            total_inference_time = 0.0
+            total_detections = 0
+            sum_confidences = 0.0
+            total_conf_count = 0
+
+            for img_file in images:
+                # Update progress/ETA before processing next image
+                steps_done += 1
+                fraction_done = steps_done / total_inferences
+                progress_bar.progress(fraction_done)
+
+                elapsed_time = time.time() - start_time
+                time_per_step = elapsed_time / steps_done
+                remaining_steps = total_inferences - steps_done
+                eta_seconds = remaining_steps * time_per_step
+                eta_placeholder.info(f"Progress: {fraction_done:.1%}. ETA: ~{eta_seconds:.1f} s")
+
+                # Load image
+                try:
+                    pil_img = Image.open(img_file).convert("RGB")
+                    np_img = np.array(pil_img)
+                except Exception as e:
+                    st.error(f"Error reading image {img_file.name}: {e}")
+                    continue
+
+                # Run inference
+                try:
+                    result = model(np_img)
+                except Exception as e:
+                    st.error(f"Inference error for model {model_file.name} on {img_file.name}: {e}")
+                    continue
+
+                r = result[0]
+                model_image_results[model_file.name][img_file.name] = r
+
+                # Accumulate inference time
+                if isinstance(r.speed, dict) and "inference" in r.speed:
+                    total_inference_time += r.speed["inference"]
+
+                # Count detections & confidence
+                if r.boxes is not None:
+                    det_count = len(r.boxes)
+                    total_detections += det_count
+                    if det_count > 0:
+                        confs = r.boxes.conf.cpu().numpy() if hasattr(r.boxes.conf, "cpu") else r.boxes.conf
+                        sum_confidences += confs.sum()
+                        total_conf_count += det_count
+
+            # After all images for this model
+            image_count = len(images)
+            avg_inference_time = total_inference_time / image_count if image_count > 0 else float("inf")
+            avg_confidence = sum_confidences / total_conf_count if total_conf_count > 0 else 0.0
+
+            model_agg_data[model_file.name] = {
+                "Model File": model_file.name,
+                "Avg Inference Time (ms)": avg_inference_time,
+                "Total Detections": total_detections,
+                "Average Confidence": avg_confidence
+            }
+
+        if not model_agg_data:
+            st.write("No valid models processed.")
+            return
+
+        # Now that all inferences are done, remove the ETA info
+        eta_placeholder.empty()
+
+        # Display aggregated metrics
+        df = pd.DataFrame(model_agg_data.values())
+        st.subheader("Aggregated Metrics (Across All Images)")
+        st.dataframe(df, use_container_width=True)
+
+        # Weighted Scoring with reciprocal-based speed
+        detection_max = df["Total Detections"].max()
+        confidence_max = df["Average Confidence"].max()
+        if detection_max == 0: detection_max = 1
+        if confidence_max == 0: confidence_max = 1
+
+        df["Detection Norm"] = df["Total Detections"] / detection_max
+        df["Confidence Norm"] = df["Average Confidence"] / confidence_max
+
+        # Convert time to speed = 1 / time, then normalize
+        eps = 1e-9
+        df["Speed Val"] = 1.0 / (df["Avg Inference Time (ms)"] + eps)
+        max_speed_val = df["Speed Val"].max() if not df["Speed Val"].isnull().all() else 1
+        if max_speed_val == 0:
+            max_speed_val = 1
+
+        df["Speed Norm"] = df["Speed Val"] / max_speed_val
+
+        df["Weighted Score"] = (
+            alpha_detection * df["Detection Norm"] +
+            beta_confidence * df["Confidence Norm"] +
+            gamma_speed     * df["Speed Norm"]
+        )
+
+        st.subheader("Weighted Score Analysis")
+        st.write(f"Weights: Detection={alpha_detection}, Confidence={beta_confidence}, Speed={gamma_speed}")
+        st.dataframe(df[[
+            "Model File",
+            "Avg Inference Time (ms)",
+            "Total Detections",
+            "Average Confidence",
+            "Detection Norm",
+            "Confidence Norm",
+            "Speed Val",
+            "Speed Norm",
+            "Weighted Score"
+        ]], use_container_width=True)
+
+        # Identify best overall model (highest Weighted Score)
+        best_idx = df["Weighted Score"].idxmax()
+        best_model = df.loc[best_idx, "Model File"]
+        best_score = df.loc[best_idx, "Weighted Score"]
+
+        st.markdown(f"""
+        **Best Overall Model** based on Weighted Score:
+        **{best_model}** (Score: {best_score:.3f}).
+
+        ### Interpretation:
+        - **Detection Norm** → fraction of the best detection count.
+        - **Confidence Norm** → fraction of the highest average confidence.
+        - **Speed Norm** → fraction of the highest (1/time). The fastest model is near 1; others are a fraction of that speed.
+
+        If you find one factor more important, adjust the weights:
+        - Increase **Detection** weight if you care about finding as many objects as possible.
+        - Increase **Confidence** weight if you only trust high‐confidence detections.
+        - Increase **Speed** weight if you need real‐time inference.
+        """)
+
+        # Display annotated images in a grid (row = image, column = model)
+        st.subheader("Annotated Images Grid (Row = Image, Column = Model)")
+        model_names_sorted = sorted(model_agg_data.keys())
+
+        for img_file in images:
+            st.markdown(f"### Image: {img_file.name}")
+            columns = st.columns(len(model_names_sorted))
+            for col, model_name in zip(columns, model_names_sorted):
+                r = model_image_results.get(model_name, {}).get(img_file.name, None)
+                if r is None:
+                    col.write(f"No results for {model_name}")
+                    continue
+
+                try:
+                    # Return a PIL image in correct RGB color space
+                    annotated_img_pil = r.plot(conf=True, boxes=True, labels=True, pil=True)
+                    col.image(
+                        annotated_img_pil,
+                        caption=f"{model_name}",
+                        use_container_width=True
+                    )
+                except Exception as e:
+                    col.error(f"Error annotating image for {model_name}: {e}")
+
+def main():
+    st.sidebar.title("Navigation")
+    page = st.sidebar.radio("Go to", ("YOLO Model Inference Tool", "YOLO Models Comparison Tool"))
+
+    if page == "YOLO Model Inference Tool":
+        yolo_inference_tool()
+    else:
+        yolo_model_comparison_tool()
+
+if __name__ == "__main__":
+    main()