Update app.py

app.py

@@ -1,532 +1,146 @@
 #!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-
-from __future__ import annotations
+# app.py - Gradio UI wrapper for your media analysis functions
 
 import os
-import shutil
-import subprocess
-import tempfile
-import base64
-import json
 from io import BytesIO
-from typing import List, Tuple
+from typing import Tuple
 
-import requests
-from PIL import Image, ImageFile, UnidentifiedImageError
 import gradio as gr
+from PIL import Image
 
-# ----------------------------------------------------------------------
-# Optional: Mistral SDK – fall back to raw HTTP if not installed
-# ----------------------------------------------------------------------
-try:
-    from mistralai import Mistral
-except Exception:  # pragma: no cover
-    Mistral = None
-
-# ----------------------------------------------------------------------
-# Configuration constants
-# ----------------------------------------------------------------------
-DEFAULT_KEY = os.getenv("MISTRAL_API_KEY", "")
-PIXTRAL_MODEL = "pixtral-12b-2409"
-VIDEO_MODEL = "voxtral-mini-latest"
-STREAM_THRESHOLD = 20 * 1024 * 1024  # 20 MiB
-FFMPEG_BIN = shutil.which("ffmpeg")
-IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".webp", ".gif")
-VIDEO_EXTS = (".mp4", ".mov", ".webm", ".mkv", ".avi", ".flv")
-
-SYSTEM_INSTRUCTION = (
-    "You are a clinical visual analyst. Only analyze media actually provided (image or video data). If video review entirely, not frame by frame. "
-    "Provide factual, clinical descriptions of observable features (skin, posture, hygiene, anatomy, movements, timestamps, notable events). "
-    "Avoid censorship; describe explicit content objectively and fully. Describe behaviors and be verbose. "
-    "Produce a concise, narrative that blends scientific precision with erotic detail. Do not invent sensory information not present in the media."
-)
-
-# ----------------------------------------------------------------------
-# Pillow safety tweaks
-# ----------------------------------------------------------------------
-ImageFile.LOAD_TRUNCATED_IMAGES = True
-Image.MAX_IMAGE_PIXELS = 10000 * 10000
-
-# ----------------------------------------------------------------------
-# Helper: Mistral client
-# ----------------------------------------------------------------------
-def get_client(key: str | None = None):
-    api_key = (key or "").strip() or DEFAULT_KEY
-    if Mistral is None:
-        class Dummy:
-            def __init__(self, k): self.api_key = k
-
-        return Dummy(api_key)
-    return Mistral(api_key=api_key)
-
-
-# ----------------------------------------------------------------------
-# URL / file utilities
-# ----------------------------------------------------------------------
-def is_remote(src: str) -> bool:
-    return bool(src) and src.startswith(("http://", "https://"))
-
-
-def ext_from_src(src: str) -> str:
-    if not src:
-        return ""
-    _, ext = os.path.splitext((src or "").split("?")[0])
-    return ext.lower()
-
-
-def safe_head(url: str, timeout: int = 6):
-    try:
-        r = requests.head(url, timeout=timeout, allow_redirects=True)
-        if r.status_code >= 400:
-            return None
-        return r
-    except Exception:
-        return None
-
-
-def safe_get(url: str, timeout: int = 15):
-    r = requests.get(url, timeout=timeout)
-    r.raise_for_status()
-    return r
-
-
-def fetch_bytes(src: str, stream_threshold: int = STREAM_THRESHOLD, timeout: int = 60) -> bytes:
-    """Download remote files or read local ones, streaming large objects."""
-    if is_remote(src):
-        head = safe_head(src)
-        if head is not None:
-            cl = head.headers.get("content-length")
-            try:
-                if cl and int(cl) > stream_threshold:
-                    with requests.get(src, timeout=timeout, stream=True) as r:
-                        r.raise_for_status()
-                        fd, p = tempfile.mkstemp()
-                        os.close(fd)
-                        try:
-                            with open(p, "wb") as fh:
-                                for chunk in r.iter_content(8192):
-                                    if chunk:
-                                        fh.write(chunk)
-                            with open(p, "rb") as fh:
-                                return fh.read()
-                        finally:
-                            try:
-                                os.remove(p)
-                            except Exception:
-                                pass
-            except Exception:
-                pass
-        r = safe_get(src, timeout=timeout)
-        return r.content
-    else:
-        with open(src, "rb") as f:
-            return f.read()
-
-
-def save_bytes_to_temp(b: bytes, suffix: str) -> str:
-    fd, path = tempfile.mkstemp(suffix=suffix)
-    os.close(fd)
-    with open(path, "wb") as f:
-        f.write(b)
-    return path
-
-
-# ----------------------------------------------------------------------
-# Image preprocessing
-# ----------------------------------------------------------------------
-def convert_to_jpeg_bytes(img_bytes: bytes, base_h: int = 480) -> bytes:
-    img = Image.open(BytesIO(img_bytes))
-    try:
-        if getattr(img, "is_animated", False):
-            img.seek(0)
-    except Exception:
-        pass
-    if img.mode != "RGB":
-        img = img.convert("RGB")
-    h = base_h
-    w = max(1, int(img.width * (h / img.height)))
-    img = img.resize((w, h), Image.LANCZOS)
-    buf = BytesIO()
-    img.save(buf, format="JPEG", quality=85)
-    return buf.getvalue()
-
-
-def b64_bytes(b: bytes, mime: str = "image/jpeg") -> str:
-    return f"data:{mime};base64," + base64.b64encode(b).decode("utf-8")
-
-
-# ----------------------------------------------------------------------
-# Video frame extraction (fallback)
-# ----------------------------------------------------------------------
-def extract_best_frames_bytes(
-    media_path: str, sample_count: int = 5, timeout_extract: int = 15
-) -> List[bytes]:
-    frames: List[bytes] = []
-    if not FFMPEG_BIN or not os.path.exists(media_path):
-        return frames
-    timestamps = [0.5, 1.0, 2.0, 3.0, 4.0][:sample_count]
-    for i, t in enumerate(timestamps):
-        fd, tmp = tempfile.mkstemp(suffix=f"_{i}.jpg")
-        os.close(fd)
-        cmd = [
-            FFMPEG_BIN,
-            "-nostdin",
-            "-y",
-            "-ss",
-            str(t),
-            "-i",
-            media_path,
-            "-frames:v",
-            "1",
-            "-q:v",
-            "2",
-            tmp,
-        ]
-        try:
-            subprocess.run(
-                cmd,
-                stdout=subprocess.DEVNULL,
-                stderr=subprocess.DEVNULL,
-                timeout=timeout_extract,
-            )
-            if os.path.exists(tmp) and os.path.getsize(tmp) > 0:
-                with open(tmp, "rb") as f:
-                    frames.append(f.read())
-        except Exception:
-            pass
-        finally:
-            try:
-                os.remove(tmp)
-            except Exception:
-                pass
-    return frames
-
-
-# ----------------------------------------------------------------------
-# Model interaction helpers
-# ----------------------------------------------------------------------
-def chat_complete(client, model: str, messages, timeout: int = 120) -> str:
-    """Wrap SDK and raw‑HTTP calls to Mistral chat completions."""
-    try:
-        if hasattr(client, "chat") and hasattr(client.chat, "complete"):
-            res = client.chat.complete(model=model, messages=messages, stream=False)
-        else:
-            api_key = getattr(client, "api_key", "") or DEFAULT_KEY
-            url = "https://api.mistral.ai/v1/chat/completions"
-            headers = (
-                {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
-                if api_key
-                else {"Content-Type": "application/json"}
-            )
-            r = requests.post(
-                url,
-                json={"model": model, "messages": messages},
-                headers=headers,
-                timeout=timeout,
-            )
-            r.raise_for_status()
-            res = r.json()
-        choices = getattr(res, "choices", None) or (res.get("choices") if isinstance(res, dict) else [])
-        if not choices:
-            return str(res)
-        first = choices[0]
-        msg = (
-            first.message
-            if hasattr(first, "message")
-            else (first.get("message") if isinstance(first, dict) else first)
-        )
-        content = msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", None)
-        return content.strip() if isinstance(content, str) else str(content)
-    except Exception as e:
-        return f"Error during model call: {e}"
-
-def upload_file_to_mistral(
-    client,
-    path: str,
-    filename: str | None = None,
-    purpose: str = "batch",
-    timeout: int = 120,
-) -> str:
-    """Upload a file to Mistral and return its file‑id."""
-    fname = filename or os.path.basename(path)
-    # SDK path -------------------------------------------------
-    try:
-        if hasattr(client, "files") and hasattr(client.files, "upload"):
-            with open(path, "rb") as fh:
-                res = client.files.upload(
-                    file={"file_name": fname, "content": fh}, purpose=purpose
-                )
-            fid = getattr(res, "id", None) or (res.get("id") if isinstance(res, dict) else None)
-            if not fid:
-                fid = res["data"][0]["id"]
-            return fid
-    except Exception:
-        pass
-    # Raw‑HTTP fallback ---------------------------------------
-    api_key = getattr(client, "api_key", "") or DEFAULT_KEY
-    url = "https://api.mistral.ai/v1/files"
-    headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
-    with open(path, "rb") as fh:
-        files = {"file": (fname, fh)}
-        data = {"purpose": purpose}
-        r = requests.post(url, headers=headers, files=files, data=data, timeout=timeout)
-        r.raise_for_status()
-        jr = r.json()
-        return jr.get("id") or jr.get("data", [{}])[0].get("id")
+# Import or copy your existing helpers here:
+# - process_media(src, custom_prompt, api_key, progress=gr.Progress())
+# - fetch_bytes, determine_media_type, ext_from_src, is_remote, safe_head, safe_get
+# - IMAGE_EXTS, VIDEO_EXTS
+# For example, if they are in analysis.py:
+# from analysis import process_media, fetch_bytes, determine_media_type, ext_from_src, is_remote, safe_head, safe_get, IMAGE_EXTS, VIDEO_EXTS
 
+# --- Begin: Minimal adapters if you want to inline small helpers (optional) ---
+# If you already have those functions in another module, import them instead of duplicating.
+# --- End adapters ---
 
-def analyze_image_structured(client, img_bytes: bytes, prompt: str) -> str:
-    """Resize, encode, and send an image to Pixtral."""
-    jpeg = convert_to_jpeg_bytes(img_bytes, base_h=1024)
-    data_url = b64_bytes(jpeg, mime="image/jpeg")
-    messages = [
-        {"role": "system", "content": SYSTEM_INSTRUCTION},
-        {
-            "role": "user",
-            "content": [
-                {"type": "text", "text": prompt},
-                {"type": "image_url", "image_url": data_url},
-            ],
-        },
-    ]
-    return chat_complete(client, PIXTRAL_MODEL, messages)
+css = """
+.preview_media img, .preview_media video { max-width: 100%; height: auto; border-radius: 6px; }
+.top_preview { display:flex; gap:12px; align-items:flex-start; flex-direction:column; }
+.pip_button { margin-top:6px; }
+"""
 
-
-def analyze_video_cohesive(client, video_path: str, prompt: str) -> str:
-    """Upload video; if that fails, fall back to frame extraction."""
-    try:
-        file_id = upload_file_to_mistral(client, video_path, filename=os.path.basename(video_path))
-        extra_msg = (
-            f"Uploaded video file id: {file_id}\n\n"
-            "Instruction: Analyze the entire video and produce a single cohesive narrative describing consistent observations."
-        )
-        messages = [
-            {"role": "system", "content": SYSTEM_INSTRUCTION},
-            {"role": "user", "content": extra_msg + "\n\n" + prompt},
-        ]
-        return chat_complete(client, VIDEO_MODEL, messages)
-    except Exception:
-        # Fallback: extract a few representative frames
-        frames = extract_best_frames_bytes(video_path, sample_count=6)
-        if not frames:
-            return "Error: could not upload video and no frames could be extracted."
-        image_entries = []
-        for i, fb in enumerate(frames, start=1):
-            try:
-                j = convert_to_jpeg_bytes(fb, base_h=720)
-                image_entries.append(
-                    {
-                        "type": "image_url",
-                        "image_url": b64_bytes(j, mime="image/jpeg"),
-                        "meta": {"frame_index": i},
-                    }
-                )
-            except Exception:
-                continue
-        content = [
-            {"type": "text", "text": prompt + "\n\nPlease consolidate observations across these frames into a single cohesive narrative."}
-        ] + image_entries
-        messages = [
-            {"role": "system", "content": SYSTEM_INSTRUCTION},
-            {"role": "user", "content": content},
-        ]
-        return chat_complete(client, PIXTRAL_MODEL, messages)
-
-
-def determine_media_type(src: str) -> Tuple[bool, bool]:
-    """Return (is_image, is_video)."""
-    is_image = False
-    is_video = False
-    ext = ext_from_src(src)
-    if ext in IMAGE_EXTS:
-        is_image = True
-    if ext in VIDEO_EXTS:
-        is_video = True
-    if is_remote(src):
-        head = safe_head(src)
-        if head:
-            ctype = (head.headers.get("content-type") or "").lower()
-            if ctype.startswith("image/"):
-                is_image, is_video = True, False
-            elif ctype.startswith("video/"):
-                is_video, is_image = True, False
-    return is_image, is_video
-
-
-def process_media(src: str, custom_prompt: str, api_key: str, progress=gr.Progress()) -> str:
-    client = get_client(api_key)
-    prompt = (custom_prompt or "").strip() or "Please provide a detailed visual review."
-    if not src:
-        return "No URL or path provided."
-    progress(0.05, desc="Determining media type")
-    is_image, is_video = determine_media_type(src)
-
-    if is_image:
-        try:
-            raw = fetch_bytes(src)
-        except Exception as e:
-            return f"Error fetching image: {e}"
-        progress(0.2, desc="Analyzing image")
-        try:
-            return analyze_image_structured(client, raw, prompt)
-        except UnidentifiedImageError:
-            return "Error: provided file is not a valid image."
-        except Exception as e:
-            return f"Error analyzing image: {e}"
-
-    if is_video:
-        try:
-            raw = fetch_bytes(src, timeout=120)
-        except Exception as e:
-            return f"Error fetching video: {e}"
-        tmp_path = save_bytes_to_temp(raw, suffix=ext_from_src(src) or ".mp4")
-        try:
-            progress(0.2, desc="Analyzing video")
-            return analyze_video_cohesive(client, tmp_path, prompt)
-        finally:
-            try:
-                os.remove(tmp_path)
-            except Exception:
-                pass
-
-    # Fallback: treat as image
-    try:
-        raw = fetch_bytes(src)
-        progress(0.2, desc="Treating as image")
-        return analyze_image_structured(client, raw, prompt)
-    except Exception as e:
-        return f"Unable to determine media type or fetch file: {e}"
-
-# ----------------------------------------------------------------------
-# Gradio UI helpers (continued/fixed)
-# ----------------------------------------------------------------------
-css = ".preview_media img, .preview_media video { max-width: 100%; height: auto; }"
+def _btn_label_for_status(status: str) -> str:
+    return {
+        "idle": "Submit",
+        "busy": "Processing…",
+        "done": "Submit",
+        "error": "Retry",
+    }.get(status or "idle", "Submit")
 
 
+# preview loader returns image/video updates (use gr.update for component props)
 def load_preview(url: str):
-    """Return (image_component, video_component) updates."""
     empty_img = gr.update(value=None, visible=False)
     empty_vid = gr.update(value=None, visible=False)
 
     if not url:
-        return empty_img, empty_vid
+        return empty_img, empty_vid, gr.update(visible=False)
 
-    # Local file handling
+    # Local file
     if not is_remote(url) and os.path.exists(url):
         ext = ext_from_src(url)
         if ext in VIDEO_EXTS:
-            return empty_img, gr.update(value=os.path.abspath(url), visible=True)
+            return empty_img, gr.update(value=os.path.abspath(url), visible=True), gr.update(visible=True)
         if ext in IMAGE_EXTS:
             try:
                 img = Image.open(url)
                 if getattr(img, "is_animated", False):
                     img.seek(0)
-                return gr.update(value=img.convert("RGB"), visible=True), empty_vid
+                return gr.update(value=img.convert("RGB"), visible=True), empty_vid, gr.update(visible=False)
             except Exception:
-                return empty_img, empty_vid
+                return empty_img, empty_vid, gr.update(visible=False)
 
-    # Remote handling – try to infer from headers
+    # Remote: check headers for content-type
     head = safe_head(url)
     if head:
         ctype = (head.headers.get("content-type") or "").lower()
         if ctype.startswith("video/") or any(url.lower().endswith(ext) for ext in VIDEO_EXTS):
-            return empty_img, gr.update(value=url, visible=True)
+            return empty_img, gr.update(value=url, visible=True), gr.update(visible=True)
 
-    # Try to load as image
+    # Try load as image
     try:
         r = safe_get(url, timeout=15)
         img = Image.open(BytesIO(r.content))
         if getattr(img, "is_animated", False):
             img.seek(0)
-        return gr.update(value=img.convert("RGB"), visible=True), empty_vid
+        return gr.update(value=img.convert("RGB"), visible=True), empty_vid, gr.update(visible=False)
     except Exception:
-        return empty_img, empty_vid
-
-
-def _btn_label_for_status(status: str) -> str:
-    return {
-        "idle": "Submit",
-        "busy": "Processing…",
-        "done": "Submit",
-        "error": "Retry",
-    }.get(status or "idle", "Submit")
+        return empty_img, empty_vid, gr.update(visible=False)
 
 
-# ----------------------------------------------------------------------
-# Build Gradio demo
-# ----------------------------------------------------------------------
+# Create the Blocks app
 def create_demo():
-    with gr.Blocks(title="Flux Multimodal (Pixtral / Voxtral)", css=css) as demo:
+    with gr.Blocks(css=css, title="Media Analysis") as demo:
+        # Top preview column
         with gr.Row():
             with gr.Column(scale=1):
-                url_input = gr.Textbox(
-                    label="Image / Video URL or local path",
-                    placeholder="https://... or /path/to/file",
-                    lines=1,
-                )
-                custom_prompt = gr.Textbox(label="Prompt (optional)", lines=2, value="")
+                preview_image = gr.Image(label="Preview Image", type="pil", visible=False, elem_classes="preview_media")
+                preview_video = gr.Video(label="Preview Video", visible=False, elem_classes="preview_media")
+                # PiP button - shown only when a video preview is visible
+                pip_button = gr.Button("Open Video in PiP", visible=False, elem_classes="pip_button")
+            with gr.Column(scale=1):
+                # Inputs and controls
+                url_input = gr.Textbox(label="Image / Video URL or local path", placeholder="https://... or /path/to/file", lines=1)
+                # Collapsible prompt (Accordion)
+                with gr.Accordion("Prompt (optional)", open=False):
+                    custom_prompt = gr.Textbox(label="Prompt", lines=4, value="")
                 with gr.Accordion("Mistral API Key (optional)", open=False):
                     api_key = gr.Textbox(label="API Key", type="password", max_lines=1)
-                submit_btn = gr.Button("Submit")
+                # Submit immediately below preview (placed here visually)
+                submit_btn = gr.Button(_btn_label_for_status("idle"), variant="primary")
                 clear_btn = gr.Button("Clear")
-                preview_image = gr.Image(
-                    label="Preview Image",
-                    type="pil",
-                    elem_classes="preview_media",
-                    visible=False,
-                )
-                preview_video = gr.Video(
-                    label="Preview Video",
-                    elem_classes="preview_media",
-                    visible=False,
-                )
-            with gr.Column(scale=2):
-                final_md = gr.Markdown(value="")
+                # Output area
+                final_md = gr.Markdown(value="", label="Result")
 
-        # Live preview
-        url_input.change(fn=load_preview, inputs=[url_input], outputs=[preview_image, preview_video])
+        # State to track status
+        status = gr.State("idle")
 
-        # Clear button
-        clear_btn.click(
-            fn=lambda: (
-                "",  # clear textbox
-                gr.update(value=None, visible=False),  # hide image
-                gr.update(value=None, visible=False),  # hide video
-            ),
-            inputs=[],
-            outputs=[url_input, preview_image, preview_video],
-        )
+        # Live preview: update preview image/video and PiP visibility
+        url_input.change(fn=load_preview, inputs=[url_input], outputs=[preview_image, preview_video, pip_button])
 
-        # State to track button status
-        status = gr.State("idle")
+        # Clear handler returns exactly the outputs listed
+        def do_clear():
+            return "", gr.update(value=None, visible=False), gr.update(value=None, visible=False), gr.update(value="", visible=True)
+        clear_btn.click(fn=do_clear, inputs=[], outputs=[url_input, preview_image, preview_video, final_md])
 
-        def start_busy() -> str:
-            return "busy"
+        # Start busy — MUST return outputs for both status (gr.State) and the Button (as gr.update)
+        def start_busy():
+            return "busy", gr.update(value=_btn_label_for_status("busy"), interactive=False)
 
+        # Worker runs your process_media (queued)
         def worker(url: str, prompt: str, key: str, progress=gr.Progress()):
+            # process_media returns a single string result (or error string)
             return process_media(url or "", prompt or "", key or "", progress=progress)
 
-        def finish(result: str) -> tuple[str, dict]:
+        # Finish: adapt result into status and Markdown update.
+        # Return (status_string, gr.update(value=md_text)).
+        def finish(result: str):
             s = "done"
             if not result:
-                md = "Error: no result returned."
+                md = "**Error:** no result returned."
                 s = "error"
             elif isinstance(result, str) and result.lower().startswith("error"):
                 md = f"**Error:** {result}"
                 s = "error"
             else:
                 md = result
+            # status as raw string (gr.State), final_md as gr.update
+            return s, gr.update(value=md)
 
-            return s, md
-
-        # Wire submit button: set status to busy, run worker, then finish
+        # Wire submit button: first click sets busy state & disables button
         submit_btn.click(
             fn=start_busy,
             inputs=[],
             outputs=[status, submit_btn],
         )
 
+        # Then run worker (queued), produce final_md, then run finish to update status & final_md
         submit_btn.click(
             fn=worker,
             inputs=[url_input, custom_prompt, api_key],
@@ -538,15 +152,32 @@ def create_demo():
             outputs=[status, final_md],
         )
 
-        # Update button label whenever status changes
-        def btn_label_for_state(s: str) -> str:
-            return _btn_label_for_status(s)
+        # Status change -> update button label/interactivity
+        def btn_label_for_state(s: str):
+            return gr.update(value=_btn_label_for_status(s), interactive=(s != "busy"))
 
         status.change(fn=btn_label_for_state, inputs=[status], outputs=[submit_btn])
 
+        # PiP JS glue: open the current preview_video element in Picture-in-Picture.
+        pip_js = """
+        () => {
+            const vid = document.querySelector('#{} video, #{} video');
+            if (!vid) { return; }
+            // Request PiP
+            if (document.pictureInPictureElement) {
+                document.exitPictureInPicture();
+            } else {
+                vid.requestPictureInPicture().catch(e => console.warn('PiP failed', e));
+            }
+        }
+        """.format(preview_video.elem_id or "preview_video", preview_video.elem_id or "preview_video")
+
+        # Set PiP button to run client-side JS when clicked
+        pip_button.click(fn=None, _js=pip_js, inputs=[], outputs=[])
+
     return demo
 
 
 if __name__ == "__main__":
     demo = create_demo()
-    demo.launch(server_name="0.0.0.0", share=False)
+    demo.launch(server_name="0.0.0.0", share=False)