Upload Imgscope-OCR-2B-0527 [ Video Understanding ] Demo (#1)

- Upload Imgscope-OCR-2B-0527 [ Video Understanding ] Demo (d535976fa172999d3122fe6636d9a0764b8e48ac)

Imgscope-OCR-2B-05270-Video-Understanding/Imgscope-OCR-2B-0527-Video-Understanding.ipynb

@@ -0,0 +1,164 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "gpuType": "T4"
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": 1,
+      "metadata": {
+        "id": "XKQwuI75LWLA"
+      },
+      "outputs": [],
+      "source": [
+        "%%capture\n",
+        "!pip install gradio transformers pillow opencv-python\n",
+        "!pip install accelerate torchvision torch huggingface_hub\n",
+        "!pip install hf_xet qwen-vl-utils gradio_client\n",
+        "!pip install transformers-stream-generator spaces"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import os\n",
+        "import uuid\n",
+        "import time\n",
+        "from threading import Thread\n",
+        "\n",
+        "import gradio as gr\n",
+        "import torch\n",
+        "import numpy as np\n",
+        "import cv2\n",
+        "from PIL import Image\n",
+        "from transformers import Qwen2VLForConditionalGeneration, AutoProcessor\n",
+        "\n",
+        "# Ensure CUDA if available\n",
+        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
+        "\n",
+        "# Load Callisto OCR3 multimodal model and processor\n",
+        "MODEL_ID = \"prithivMLmods/Imgscope-OCR-2B-0527\"\n",
+        "processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)\n",
+        "model = Qwen2VLForConditionalGeneration.from_pretrained(\n",
+        "    MODEL_ID,\n",
+        "    trust_remote_code=True,\n",
+        "    torch_dtype=torch.float16\n",
+        ").to(device).eval()\n",
+        "\n",
+        "# Constants\n",
+        "MAX_INPUT_TOKEN_LENGTH = 4096\n",
+        "\n",
+        "\n",
+        "def downsample_video(video_path: str, num_frames: int = 10):\n",
+        "    \"\"\"\n",
+        "    Extracts 'num_frames' evenly spaced frames from the video.\n",
+        "    Returns a list of (PIL.Image, timestamp_seconds).\n",
+        "    \"\"\"\n",
+        "    vidcap = cv2.VideoCapture(video_path)\n",
+        "    total = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
+        "    fps = vidcap.get(cv2.CAP_PROP_FPS) or 1\n",
+        "    indices = np.linspace(0, total - 1, num_frames, dtype=int)\n",
+        "    frames = []\n",
+        "    for idx in indices:\n",
+        "        vidcap.set(cv2.CAP_PROP_POS_FRAMES, idx)\n",
+        "        ret, frame = vidcap.read()\n",
+        "        if not ret:\n",
+        "            continue\n",
+        "        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n",
+        "        pil = Image.fromarray(frame)\n",
+        "        timestamp = round(idx / fps, 2)\n",
+        "        frames.append((pil, timestamp))\n",
+        "    vidcap.release()\n",
+        "    return frames\n",
+        "\n",
+        "\n",
+        "def generate(video_file: str):\n",
+        "    \"\"\"\n",
+        "    Process the uploaded video through OCR and return concatenated output.\n",
+        "    \"\"\"\n",
+        "    # Step 1: extract frames\n",
+        "    frames = downsample_video(video_file)\n",
+        "\n",
+        "    # Step 2: build chat-like messages\n",
+        "    messages = [\n",
+        "        {\"role\": \"system\", \"content\": [{\"type\": \"text\", \"text\": \"You are a helpful assistant, for video understanding.\"}]},\n",
+        "        {\"role\": \"user\", \"content\": [{\"type\": \"text\", \"text\": \"Please describe the content of the following video frames:\"}]\n",
+        "        }\n",
+        "    ]\n",
+        "    for img, ts in frames:\n",
+        "        # save temporary frame image\n",
+        "        path = f\"frame_{uuid.uuid4().hex}.png\"\n",
+        "        img.save(path)\n",
+        "        messages[1][\"content\"].append({\"type\": \"text\", \"text\": f\"Frame at {ts}s:\"})\n",
+        "        messages[1][\"content\"].append({\"type\": \"image\", \"url\": path})\n",
+        "\n",
+        "    # Step 3: tokenize with truncation\n",
+        "    inputs = processor.apply_chat_template(\n",
+        "        messages,\n",
+        "        tokenize=True,\n",
+        "        add_generation_prompt=True,\n",
+        "        return_dict=True,\n",
+        "        return_tensors=\"pt\",\n",
+        "        truncation=True,\n",
+        "        max_length=MAX_INPUT_TOKEN_LENGTH\n",
+        "    ).to(device)\n",
+        "\n",
+        "    # Step 4: use streamer to collect output\n",
+        "    from transformers import TextIteratorStreamer\n",
+        "    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)\n",
+        "    gen_kwargs = {\n",
+        "        **inputs,\n",
+        "        \"streamer\": streamer,\n",
+        "        \"max_new_tokens\": 1024,\n",
+        "        \"do_sample\": True,\n",
+        "        \"temperature\": 0.7,\n",
+        "    }\n",
+        "    thread = Thread(target=model.generate, kwargs=gen_kwargs)\n",
+        "    thread.start()\n",
+        "\n",
+        "    # collect all tokens\n",
+        "    buffer = \"\"\n",
+        "    for chunk in streamer:\n",
+        "        buffer += chunk.replace(\"<|im_end|>\", \"\")\n",
+        "        time.sleep(0.01)\n",
+        "\n",
+        "    # return full concatenated response\n",
+        "    return buffer\n",
+        "\n",
+        "\n",
+        "def launch_app():\n",
+        "    demo = gr.Interface(\n",
+        "        fn=generate,\n",
+        "        inputs=gr.Video(label=\"Upload Video\"),\n",
+        "        outputs=gr.Textbox(label=\"Video Description\"),\n",
+        "        title=\"Video Understanding with Imgscope-OCR-2B-0527\",\n",
+        "        description=\"Upload a video and get an OCR-based description of its frames.\",\n",
+        "        allow_flagging=\"never\"\n",
+        "    )\n",
+        "    demo.queue().launch(debug=True)\n",
+        "\n",
+        "\n",
+        "if __name__ == \"__main__\":\n",
+        "    launch_app()"
+      ],
+      "metadata": {
+        "id": "GZXqC00zLbS1"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}