MCPcopy Create free account
hub / github.com/LargeWorldModel/LWM / _read_process_vision

Method _read_process_vision

lwm/vision_chat.py:76–108  ·  view source on GitHub ↗
(self, path, max_n_frames)

Source from the content-addressed store, hash-verified

74 return np.array(image, dtype=np.float32) / 127.5 - 1
75
76 def _read_process_vision(self, path, max_n_frames):
77 f = open_file(path, 'rb')
78 if path.endswith('.png') or path.endswith('.jpg'):
79 image = Image.open(f).convert('RGB')
80 vision = self._process_frame(image, 256)[None]
81 else:
82 vr = decord.VideoReader(f, ctx=decord.cpu(0))
83 duration = len(vr)
84 if duration <= max_n_frames:
85 frame_id_list = list(range(duration))
86 else:
87 frame_id_list = np.linspace(0, duration - 1, max_n_frames, dtype=int).tolist()
88 video = vr.get_batch(frame_id_list).asnumpy()
89 vision = np.stack([self._process_frame(Image.fromarray(frame), 256) for frame in video])
90
91 B = 1
92 encodings = []
93 for i in range(0, len(vision), 1):
94 v = vision[i:i + B]
95 if len(v) % B == 0:
96 n_pad = 0
97 else:
98 n_pad = B - len(v) % B
99 v = np.pad(v, ((n_pad, 0), (0, 0), (0, 0), (0, 0)))
100 enc = jax.device_get(self.vqgan.encode(v))[1].astype(int)
101 enc = enc[n_pad:]
102 for t in range(len(enc)):
103 encodings.extend(enc[t].reshape(-1).tolist())
104 if t == len(enc) - 1:
105 encodings.append(8193)
106 else:
107 encodings.append(8192)
108 return encodings
109
110 def construct_input(self, prompts, max_n_frames):
111 max_input_length = max_n_frames * self.n_tokens_per_frame + self.min_buffer_size

Callers 1

construct_inputMethod · 0.95

Calls 2

_process_frameMethod · 0.95
encodeMethod · 0.45

Tested by

no test coverage detected