Update app.py

app.py

@@ -10,10 +10,7 @@ import uuid
 from object_detection import detectObjects, detectVideo
 from object_detection_count import detectObjectsAndCount
 from pose_analysis import process_gif
-
-# Traffic sign detection (image + video)
-from traffic_sign_detection import detectTrafficObjects, detectTrafficVideo
-
+from traffic_sign_detection import detectTrafficObjects
 
 # -----------------------------
 # Constants
@@ -33,23 +30,6 @@ TASK_TO_ASSET_SUBDIR = {
 IMAGE_EXTS = {".jpg", ".jpeg", ".png"}
 VIDEO_EXTS = {".mp4", ".mov", ".avi", ".gif"}
 
-# Per-tab allowed types for uploader + validation
-TAB_ALLOWED_EXTS = {
-    "Object Detection": IMAGE_EXTS | VIDEO_EXTS,          # image + video
-    "Traffic Sign Detection": IMAGE_EXTS | VIDEO_EXTS,    # image + video
-    "Pose Analysis": VIDEO_EXTS,                          # video only
-    "Object Counting": IMAGE_EXTS,                        # image only
-}
-
-# Streamlit uploader "type" expects extensions WITHOUT dots
-TAB_UPLOADER_TYPES = {
-    "Object Detection": ["jpg", "jpeg", "png", "gif", "mp4", "avi", "mov"],
-    "Traffic Sign Detection": ["jpg", "jpeg", "png", "gif", "mp4", "avi", "mov"],
-    "Pose Analysis": ["gif", "mp4", "avi", "mov"],
-    "Object Counting": ["jpg", "jpeg", "png"],
-}
-
-
 # -----------------------------
 # Helpers
 # -----------------------------
@@ -61,19 +41,15 @@ def check_file_size(file):
 
 
 def save_uploaded_file_to_temp(uploaded_file) -> str:
-    """
-    Save upload to a temp folder to avoid cluttering project root
-    and reduce filename collisions.
-    """
+    # NOTE: left exactly as you had it (do not change)
     safe_name = uploaded_file.name.replace("/", "_").replace("\\", "_")
-    temp_dir = Path(tempfile.gettempdir()) / "streamlit_uploads"
-    temp_dir.mkdir(parents=True, exist_ok=True)
+    save_path = "./"
 
-    out_path = temp_dir / f"{uuid.uuid4().hex}_{safe_name}"
-    with open(out_path, "wb") as f:
+    with open(safe_name, "wb") as f:
         f.write(uploaded_file.getbuffer())
 
-    return str(out_path)
+    print("Saved uploaded file to:", safe_name)
+    return str(safe_name)
 
 
 def list_demo_files(task_name: str, limit: int = 6):
@@ -111,10 +87,14 @@ def is_h264(path: str) -> bool:
         out = subprocess.check_output(
             [
                 "ffprobe",
-                "-v", "error",
-                "-select_streams", "v:0",
-                "-show_entries", "stream=codec_name",
-                "-of", "default=nk=1:nw=1",
+                "-v",
+                "error",
+                "-select_streams",
+                "v:0",
+                "-show_entries",
+                "stream=codec_name",
+                "-of",
+                "default=nk=1:nw=1",
                 path,
             ],
             text=True,
@@ -131,23 +111,30 @@ def to_h264_mp4(input_path: str) -> str:
     - Uses a unique temp output so we don't overwrite assets or collide across reruns.
     """
     if shutil.which("ffmpeg") is None:
-        st.warning("ffmpeg not found. Video may not play if not H.264.")
+        st.warning("ffmpeg not found in this environment. Video may show black screen if not H.264.")
         return input_path
 
     out_path = str(Path(tempfile.gettempdir()) / f"{uuid.uuid4().hex}.mp4")
     cmd = [
         "ffmpeg",
         "-y",
-        "-i", input_path,
-        "-c:v", "libx264",
-        "-pix_fmt", "yuv420p",
-        "-movflags", "+faststart",
-        "-c:a", "aac",
-        "-b:a", "128k",
+        "-i",
+        input_path,
+        "-c:v",
+        "libx264",
+        "-pix_fmt",
+        "yuv420p",
+        "-movflags",
+        "+faststart",
+        "-c:a",
+        "aac",
+        "-b:a",
+        "128k",
         out_path,
     ]
     try:
         subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT)
+        # Sanity check: ensure output exists and is not tiny/empty
         if not Path(out_path).exists() or Path(out_path).stat().st_size < 1024:
             st.warning("Video conversion produced an empty file; showing original.")
             return input_path
@@ -162,6 +149,8 @@ def render_media(path: Path):
     if ext in IMAGE_EXTS:
         st.image(str(path), caption=path.name, use_container_width=True)
     elif ext in VIDEO_EXTS:
+        # IMPORTANT: do NOT try to "convert-and-overwrite" showcase assets.
+        # Only convert to a temp H.264 file if needed, then pass BYTES to Streamlit.
         playable = str(path)
         if not is_h264(playable):
             playable = to_h264_mp4(playable)
@@ -171,7 +160,10 @@ def render_media(path: Path):
 
 def render_showcase(tasks, per_task_limit=6):
     st.subheader("Example outputs (what to expect)")
-    st.write("These are pre-generated results so you can see the expected output before uploading.")
+    st.write(
+        "These are pre-generated results (detected/segmented/pose analyzed/traffic signs) "
+        "so you can see the expected output before uploading."
+    )
 
     for task in tasks:
         st.markdown(f"### {task}")
@@ -179,8 +171,8 @@ def render_showcase(tasks, per_task_limit=6):
 
         if not demo_files:
             st.info(
-                f"No demo files found for **{task}**. Add files under: "
-                f"`{ASSETS_DIR / TASK_TO_ASSET_SUBDIR[task]}`"
+                f"No demo files found for **{task}**. Add images/videos under: "
+                f"{ASSETS_DIR / TASK_TO_ASSET_SUBDIR[task]}"
             )
             continue
 
@@ -192,60 +184,55 @@ def render_showcase(tasks, per_task_limit=6):
         st.divider()
 
 
-def validate_file_for_tab(file_path: str, tab_name: str) -> bool:
-    ext = Path(file_path).suffix.lower()
-    allowed = TAB_ALLOWED_EXTS.get(tab_name, set())
-    return ext in allowed
-
-
 def process_file(file_path, tab_name, confidence_score, progress_placeholder, class_type):
-    """
-    Enforces per-tab constraints:
-    - Object Detection: image + video
-    - Traffic Sign Detection: image + video
-    - Pose Analysis: video only
-    - Object Counting: image only
-    """
-    if not validate_file_for_tab(file_path, tab_name):
-        allowed = ", ".join(sorted(TAB_ALLOWED_EXTS.get(tab_name, [])))
-        st.error(f"Unsupported file for **{tab_name}**. Allowed: {allowed}")
-        return None, None
-
     progress_placeholder.info(f"Processing... Please wait. (Confidence Score: {confidence_score})")
-    time.sleep(0.5)
-
-    ext = Path(file_path).suffix.lower()
+    time.sleep(1)
 
     if tab_name == "Object Detection":
-        progress_placeholder.empty()
-        if ext in IMAGE_EXTS:
+        if file_path.lower().endswith((".jpg", ".png", ".jpeg")):
+            progress_placeholder.empty()
             img = detectObjects(file_path, confidence_score)
             return img, "image"
-        else:
+
+        elif file_path.lower().endswith((".mp4", ".avi", ".mov", ".gif")):
+            progress_placeholder.empty()
             out_video_path = detectVideo(file_path, confidence_score)
             return out_video_path, "video"
 
+        progress_placeholder.empty()
+        st.error("Unsupported file format! Please upload an image or video.")
+        return None, None
+
     elif tab_name == "Object Counting":
+        if file_path.lower().endswith((".jpg", ".png", ".jpeg")):
+            progress_placeholder.empty()
+            img, count = detectObjectsAndCount(file_path, confidence_score, class_type)
+            st.info(f"Count for class '{class_type}': {count}")
+            return img, "image"
+
+        elif file_path.lower().endswith((".mp4", ".avi", ".mov", ".gif")):
+            progress_placeholder.empty()
+            out_video_path = detectVideo(file_path, confidence_score)
+            return out_video_path, "video"
+
         progress_placeholder.empty()
-        # image only
-        img, count = detectObjectsAndCount(file_path, confidence_score, class_type)
-        st.info(f"Count for class '{class_type}': {count}")
-        return img, "image"
+        st.error("Unsupported file format! Please upload an image or video.")
+        return None, None
 
     elif tab_name == "Pose Analysis":
         progress_placeholder.empty()
-        # video only
         out_video_path = process_gif(file_path, confidence_score)
         return out_video_path, "video"
 
     elif tab_name == "Traffic Sign Detection":
-        progress_placeholder.empty()
-        if ext in IMAGE_EXTS:
+        if file_path.lower().endswith((".jpg", ".png", ".jpeg")):
+            progress_placeholder.empty()
             img = detectTrafficObjects(file_path, confidence_score)
             return img, "image"
-        else:
-            out_video_path = detectTrafficVideo(file_path, confidence_score)
-            return out_video_path, "video"
+
+        progress_placeholder.empty()
+        st.error("Unsupported file format! Please upload an image.")
+        return None, None
 
     st.error("Unknown tab selection.")
     return None, None
@@ -257,23 +244,30 @@ def process_file(file_path, tab_name, confidence_score, progress_placeholder, cl
 st.set_page_config(page_title="AI Video/Image Analysis Platform", layout="wide")
 
 st.title("AI Video/Image Analysis Platform")
-st.write("Upload a file and choose a tab for analysis.")
+st.write("Upload an image or video and choose a tab for analysis.")
 
+# Tabs for different functionalities
 tabs = st.tabs(TABS)
 
 for i, tab_name in enumerate(TABS):
     with tabs[i]:
         st.header(tab_name)
 
-        uploader_label = "Upload an Image/Video"
-        if tab_name == "Pose Analysis":
-            uploader_label = "Upload a Video (Pose Analysis supports video only)"
+        # -----------------------------
+        # MINIMAL CHANGE: per-tab uploader constraints
+        # -----------------------------
+        if tab_name in ["Object Detection", "Traffic Sign Detection"]:
+            allowed_types = ["jpg", "jpeg", "png", "gif", "mp4", "avi", "mov"]  # image + video
+        elif tab_name == "Pose Analysis":
+            allowed_types = ["gif", "mp4", "avi", "mov"]  # video only
         elif tab_name == "Object Counting":
-            uploader_label = "Upload an Image (Object Counting supports images only)"
+            allowed_types = ["jpg", "jpeg", "png"]  # image only
+        else:
+            allowed_types = ["jpg", "jpeg", "png", "gif", "mp4", "avi", "mov"]
 
         uploaded_file = st.file_uploader(
-            uploader_label,
-            type=TAB_UPLOADER_TYPES[tab_name],
+            "Upload an Image/Video",
+            type=allowed_types,
             key=f"uploader_{tab_name}",
         )
 
@@ -285,14 +279,9 @@ for i, tab_name in enumerate(TABS):
 
             st.success(f"Uploaded file: {uploaded_file.name} ({file_size:.2f} MB)")
 
+            # Save to temp and use the full path for downstream processing
             file_path = save_uploaded_file_to_temp(uploaded_file)
 
-            # Validate again (hard guard)
-            if not validate_file_for_tab(file_path, tab_name):
-                allowed = ", ".join(sorted(TAB_ALLOWED_EXTS.get(tab_name, [])))
-                st.error(f"Unsupported file for **{tab_name}**. Allowed: {allowed}")
-                continue
-
             confidence_score = st.number_input(
                 "Adjust Confidence Score",
                 min_value=0.0,
@@ -312,11 +301,25 @@ for i, tab_name in enumerate(TABS):
                     key=f"class_type_{tab_name}",
                 )
 
+            safe_name = uploaded_file.name.replace("/", "_").replace("\\", "_")
+            print("Process file called with", safe_name)
+
+            # -----------------------------
+            # MINIMAL CHANGE: hard guard (enforce per-tab constraints)
+            # -----------------------------
+            ext = os.path.splitext(safe_name)[1].lower()
+            if tab_name == "Pose Analysis" and ext in (".jpg", ".jpeg", ".png"):
+                st.error("Pose Analysis supports video only. Please upload mp4/mov/avi/gif.")
+                continue
+            if tab_name == "Object Counting" and ext in (".mp4", ".mov", ".avi", ".gif"):
+                st.error("Object Counting supports images only. Please upload jpg/jpeg/png.")
+                continue
+
             if st.button(f"Process {tab_name}", key=f"process_{tab_name}"):
                 progress_placeholder = st.empty()
                 with st.spinner("Processing... Please wait."):
                     result, result_type = process_file(
-                        file_path,
+                        safe_name,
                         tab_name,
                         confidence_score,
                         progress_placeholder,
@@ -327,9 +330,11 @@ for i, tab_name in enumerate(TABS):
                     st.success(f"{tab_name} completed successfully!")
 
                     playable = result
+                    # Convert ONLY if not already H.264 (or if ffprobe missing, this will attempt conversion)
                     if not is_h264(playable):
                         playable = to_h264_mp4(playable)
 
+                    # Use bytes for reliability on Spaces
                     st_video_file(playable, fmt="video/mp4")
 
                 if result_type == "image" and result is not None: