Views
No views yet
Otter folder to make sure it has the access to otter/modeling_otter.py.1import mimetypes
2import os
3from typing import Union
4import cv2
5import requests
6import torch
7import transformers
8from PIL import Image
9import sys
10
11# make sure you can properly access the otter folder
12from otter.modeling_otter import OtterForConditionalGeneration
13
14# Disable warnings
15requests.packages.urllib3.disable_warnings()
16
17# ------------------- Utility Functions -------------------
18
19
20def get_content_type(file_path):
21 content_type, _ = mimetypes.guess_type(file_path)
22 return content_type
23
24
25# ------------------- Image and Video Handling Functions -------------------
26
27
28def extract_frames(video_path, num_frames=16):
29 video = cv2.VideoCapture(video_path)
30 total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
31 frame_step = total_frames // num_frames
32 frames = []
33
34 for i in range(num_frames):
35 video.set(cv2.CAP_PROP_POS_FRAMES, i * frame_step)
36 ret, frame = video.read()
37 if ret:
38 frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
39 frame = Image.fromarray(frame).convert("RGB")
40 frames.append(frame)
41
42 video.release()
43 return frames
44
45
46def get_image(url: str) -> Union[Image.Image, list]:
47 if "://" not in url: # Local file
48 content_type = get_content_type(url)
49 else: # Remote URL
50 content_type = requests.head(url, stream=True, verify=False).headers.get("Content-Type")
51
52 if "image" in content_type:
53 if "://" not in url: # Local file
54 return Image.open(url)
55 else: # Remote URL
56 return Image.open(requests.get(url, stream=True, verify=False).raw)
57 elif "video" in content_type:
58 video_path = "temp_video.mp4"
59 if "://" not in url: # Local file
60 video_path = url
61 else: # Remote URL
62 with open(video_path, "wb") as f:
63 f.write(requests.get(url, stream=True, verify=False).content)
64 frames = extract_frames(video_path)
65 if "://" in url: # Only remove the temporary video file if it was downloaded
66 os.remove(video_path)
67 return frames
68 else:
69 raise ValueError("Invalid content type. Expected image or video.")
70
71
72# ------------------- OTTER Prompt and Response Functions -------------------
73
74
75def get_formatted_prompt(prompt: str) -> str:
76 return f"<image>User: {prompt} GPT:<answer>"
77
78
79def get_response(input_data, prompt: str, model=None, image_processor=None, tensor_dtype=None) -> str:
80 if isinstance(input_data, Image.Image):
81 vision_x = image_processor.preprocess([input_data], return_tensors="pt")["pixel_values"].unsqueeze(1).unsqueeze(0)
82 elif isinstance(input_data, list): # list of video frames
83 vision_x = image_processor.preprocess(input_data, return_tensors="pt")["pixel_values"].unsqueeze(0).unsqueeze(0)
84 else:
85 raise ValueError("Invalid input data. Expected PIL Image or list of video frames.")
86
87 lang_x = model.text_tokenizer(
88 [
89 get_formatted_prompt(prompt),
90 ],
91 return_tensors="pt",
92 )
93
94 bad_words_id = model.text_tokenizer(["User:", "GPT1:", "GFT:", "GPT:"], add_special_tokens=False).input_ids
95 generated_text = model.generate(
96 vision_x=vision_x.to(model.device, dtype=tensor_dtype),
97 lang_x=lang_x["input_ids"].to(model.device),
98 attention_mask=lang_x["attention_mask"].to(model.device),
99 max_new_tokens=512,
100 num_beams=3,
101 no_repeat_ngram_size=3,
102 bad_words_ids=bad_words_id,
103 )
104 parsed_output = (
105 model.text_tokenizer.decode(generated_text[0])
106 .split("<answer>")[-1]
107 .lstrip()
108 .rstrip()
109 .split("<|endofchunk|>")[0]
110 .lstrip()
111 .rstrip()
112 .lstrip('"')
113 .rstrip('"')
114 )
115 return parsed_output
116
117
118# ------------------- Main Function -------------------
119load_bit = "fp32"
120if load_bit == "fp16":
121 precision = {"torch_dtype": torch.float16}
122elif load_bit == "bf16":
123 precision = {"torch_dtype": torch.bfloat16}
124elif load_bit == "fp32":
125 precision = {"torch_dtype": torch.float32}
126
127# This model version is trained on MIMIC-IT DC dataset.
128model = OtterForConditionalGeneration.from_pretrained("luodian/OTTER-9B-DenseCaption", device_map="auto", **precision)
129tensor_dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[load_bit]
130
131model.text_tokenizer.padding_side = "left"
132tokenizer = model.text_tokenizer
133image_processor = transformers.CLIPImageProcessor()
134model.eval()
135
136while True:
137 video_url = input("Enter video path: ") # Replace with the path to your video file, could be any common format.
138
139 frames_list = get_image(video_url)
140
141 while True:
142 prompts_input = input("Enter prompts: ")
143
144 if prompts_input.lower() == "quit":
145 break
146
147 print(f"\nPrompt: {prompts_input}")
148 response = get_response(frames_list, prompts_input, model, image_processor, tensor_dtype)
149 print(f"Response: {response}")
150