ma7583/ramplus

setup uv env

#!/bin/bash
# Set up RAM+ in an isolated venv with a compatible transformers version.
# recognize-anything uses APIs removed in transformers>=4.46, so we pin to 4.45.2.
set -e

VENV=".venv-ram"
WEIGHTS_DIR="models/ram_plus"
WEIGHTS_FILE="$WEIGHTS_DIR/ram_plus_swin_large_14m.pth"

# ---------------------------------------------------------------------------
# 1. Create venv
# ---------------------------------------------------------------------------
if [ ! -d "$VENV" ]; then
    echo "Creating $VENV …"
    uv venv "$VENV" --python 3.11
else
    echo "$VENV already exists"
fi

PIP="$VENV/bin/pip"

# ---------------------------------------------------------------------------
# 2. Install PyTorch (reuse system CUDA version if available)
# ---------------------------------------------------------------------------
if ! "$VENV/bin/python" -c "import torch" 2>/dev/null; then
    echo "Installing PyTorch …"
    CUDA_VER=$(python3 -c "import torch; print('cu' + torch.version.cuda.replace('.',''))" 2>/dev/null || echo "cpu")
    if [ "$CUDA_VER" = "cpu" ]; then
        uv pip install --python "$VENV/bin/python" torch torchvision
    else
        uv pip install --python "$VENV/bin/python" torch torchvision \
            --index-url "https://download.pytorch.org/whl/$CUDA_VER"
    fi
fi

# ---------------------------------------------------------------------------
# 3. Install recognize-anything + pinned transformers
# ---------------------------------------------------------------------------
if ! "$VENV/bin/python" -c "import ram" 2>/dev/null; then
    echo "Installing recognize-anything with transformers==4.45.2 …"
    uv pip install --python "$VENV/bin/python" \
        "transformers==4.45.2" \
        timm pillow scipy fairscale \
        git+https://github.com/xinyu1205/recognize-anything.git
fi

# ---------------------------------------------------------------------------
# 4. Install Streamlit in the RAM venv
# ---------------------------------------------------------------------------
if ! "$VENV/bin/python" -c "import streamlit" 2>/dev/null; then
    echo "Installing streamlit …"
    uv pip install --python "$VENV/bin/python" streamlit
fi

# ---------------------------------------------------------------------------
# 5. Download weights
# ---------------------------------------------------------------------------
mkdir -p "$WEIGHTS_DIR"
if [ ! -f "$WEIGHTS_FILE" ]; then
    echo "Downloading RAM+ weights (~3 GB) …"
    "$VENV/bin/python" -c "
from huggingface_hub import hf_hub_download
path = hf_hub_download(
    repo_id='ma7583/ramplus',
    filename='ram_plus_swin_large_14m.pth',
    local_dir='$WEIGHTS_DIR',
)
print('Downloaded to:', path)
"
else
    echo "Weights already present at $WEIGHTS_FILE"
fi

echo ""
echo "Setup complete. Run the app with:"
echo "  $VENV/bin/streamlit run tests/explore_ram.py"

run streamlit for demo

"""RAM+ Video Object Explorer — Streamlit app.

Usage:
    streamlit run tests/explore_ram.py
"""

import json
import subprocess
import tempfile
from pathlib import Path
from PIL import Image
import streamlit as st
import torch

VIDEO_DIR  = Path(".")
WEIGHTS    = Path("models/ram_plus/ram_plus_swin_large_14m.pth")
CACHE_DIR  = Path("tests/mbs/ram_tags")
IMAGE_SIZE = 384

st.set_page_config(page_title="RAM+ Explorer", layout="wide")

# ---------------------------------------------------------------------------
# Load model
# ---------------------------------------------------------------------------

@st.cache_resource(show_spinner="Loading RAM+ model…")
def load_model(weights: str):
    from ram.models import ram_plus
    import torchvision.transforms as T

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = ram_plus(pretrained=weights, image_size=IMAGE_SIZE, vit="swin_l")
    model.eval().to(device)

    transform = T.Compose([
        T.Resize((IMAGE_SIZE, IMAGE_SIZE)),
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    return model, transform, device


if not WEIGHTS.exists():
    st.error(f"Weights not found at `{WEIGHTS}`. Run `bash run_ram_setup.sh` first.")
    st.stop()

model, transform, device = load_model(str(WEIGHTS))

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def extract_frames(video_path: Path, fps: float) -> list[tuple[float, Path]]:
    tmp = Path(tempfile.mkdtemp(prefix="ram-frames-"))
    subprocess.run(
        ["ffmpeg", "-y", "-i", str(video_path), "-vf", f"fps={fps}", "-q:v", "2",
         str(tmp / "frame_%05d.jpg")],
        capture_output=True, check=True, timeout=300,
    )
    frames = sorted(tmp.glob("frame_*.jpg"))
    return [(i / fps, p) for i, p in enumerate(frames)]


@torch.inference_mode()
def tag_frame(img_path: Path, threshold: float) -> list[str]:
    img = Image.open(img_path).convert("RGB")
    tensor = transform(img).unsqueeze(0).to(device)
    model.threshold = threshold
    tags, _ = model.generate_tag(tensor)
    return [t.strip() for t in tags[0].split("|") if t.strip()]


CACHE_DIR.mkdir(parents=True, exist_ok=True)

def cache_path(video_name: str, fps: float, threshold: float) -> Path:
    return CACHE_DIR / f"{Path(video_name).stem}_{fps}fps_t{threshold:.2f}.json"


@st.cache_data(show_spinner=False)
def get_tags(video_name: str, fps: float, threshold: float, video_dir: str) -> list[dict]:
    cp = cache_path(video_name, fps, threshold)
    if cp.exists():
        return json.loads(cp.read_text())

    frames = extract_frames(Path(video_dir) / video_name, fps)
    results = []
    bar = st.progress(0, text=f"Tagging {len(frames)} frames…")
    for i, (ts, fp) in enumerate(frames):
        results.append({"timestamp": round(ts, 2), "tags": tag_frame(fp, threshold), "frame": str(fp)})
        bar.progress((i + 1) / len(frames), text=f"Frame {i+1}/{len(frames)}")
    bar.empty()

    cp.write_text(json.dumps(results, ensure_ascii=False))
    return results

# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

with st.sidebar:
    st.header("Settings")
    video_dir_input = st.text_input("Video directory", value=str(VIDEO_DIR))
    video_dir = Path(video_dir_input)

videos = sorted(p.name for p in video_dir.glob("*.mp4")) if video_dir.exists() else []
if not videos:
    st.error(f"No .mp4 files found in `{video_dir}`")
    st.stop()

with st.sidebar:
    video_name = st.selectbox(f"Video ({len(videos)} total)", videos)

col_search, col_fps, col_thresh = st.columns([4, 1, 1])
with col_search:
    tag_filter = st.text_input("Search tags", placeholder="e.g. person, shelf, phone")
with col_fps:
    fps = st.select_slider("Sampling fps", [0.25, 0.5, 1.0, 2.0], value=0.5)
with col_thresh:
    threshold = st.slider("RAM Threshold", 0.1, 0.95, 0.68, step=0.05)

cached = cache_path(video_name, fps, threshold).exists()
st.caption(f"{'✓ cached' if cached else '⚡ will run inference'}")

with st.spinner(f"Processing {video_name}…"):
    frames = get_tags(video_name, fps, threshold, str(video_dir))

# Filter by tag
if tag_filter.strip():
    needles = [n.strip().lower() for n in tag_filter.split(",") if n.strip()]
    frames = [f for f in frames if any(
        any(n in t.lower() for t in f["tags"]) for n in needles
    )]
    st.info(f"{len(frames)} frames match **{tag_filter}**")
else:
    # Show tag cloud when no filter active
    all_tags: dict[str, int] = {}
    for f in frames:
        for t in f["tags"]:
            all_tags[t] = all_tags.get(t, 0) + 1
    top = sorted(all_tags.items(), key=lambda x: -x[1])[:60]
    st.write("**All tags:** " + "  ·  ".join(f"`{t}` ({n})" for t, n in top))
    st.divider()

# Frame grid
if frames:
    cols = st.columns(4)
    for i, frame in enumerate(frames):
        with cols[i % 4]:
            st.image(frame["frame"], use_container_width=True)
            ts = frame["timestamp"]
            m, s = divmod(int(ts), 60)
            st.caption(f"**{m}:{s:02d}**  {' · '.join(frame['tags'][:5])}")
else:
    st.warning("No frames match the search.")