File size: 6,217 Bytes
c2c97b6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 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": []
}
]
} |