| 11 | from shutil import copyfile |
| 12 | |
| 13 | class YUV_Read(): |
| 14 | def __init__(self, filepath, h, w, format='yuv420', toRGB=True): |
| 15 | |
| 16 | self.h = h |
| 17 | self.w = w |
| 18 | |
| 19 | self.fp = open(filepath, 'rb') |
| 20 | |
| 21 | if format == 'yuv420': |
| 22 | self.frame_length = int(1.5 * h * w) |
| 23 | self.Y_length = h * w |
| 24 | self.Uv_length = int(0.25 * h * w) |
| 25 | else: |
| 26 | pass |
| 27 | self.toRGB = toRGB |
| 28 | |
| 29 | def read(self, offset_frame=None): |
| 30 | if not offset_frame == None: |
| 31 | self.fp.seek(offset_frame * self.frame_length, 0) |
| 32 | |
| 33 | Y = np.fromfile(self.fp, np.uint8, count=self.Y_length) |
| 34 | U = np.fromfile(self.fp, np.uint8, count=self.Uv_length) |
| 35 | V = np.fromfile(self.fp, np.uint8, count=self.Uv_length) |
| 36 | if Y.size < self.Y_length or \ |
| 37 | U.size < self.Uv_length or \ |
| 38 | V.size < self.Uv_length: |
| 39 | return None, False |
| 40 | |
| 41 | Y = np.reshape(Y, [self.w, self.h], order='F') |
| 42 | Y = np.transpose(Y) |
| 43 | |
| 44 | U = np.reshape(U, [int(self.w / 2), int(self.h / 2)], order='F') |
| 45 | U = np.transpose(U) |
| 46 | |
| 47 | V = np.reshape(V, [int(self.w / 2), int(self.h / 2)], order='F') |
| 48 | V = np.transpose(V) |
| 49 | |
| 50 | U = np.array(Image.fromarray(U).resize([self.w, self.h])) |
| 51 | V = np.array(Image.fromarray(V).resize([self.w, self.h])) |
| 52 | |
| 53 | if self.toRGB: |
| 54 | Y = Y / 255.0 |
| 55 | U = U / 255.0 - 0.5 |
| 56 | V = V / 255.0 - 0.5 |
| 57 | |
| 58 | self.YUV = np.stack((Y, U, V), axis=-1) |
| 59 | self.RGB = (255.0 * np.clip(yuv2rgb(self.YUV), 0.0, 1.0)).astype('uint8') |
| 60 | |
| 61 | self.YUV = None |
| 62 | return self.RGB, True |
| 63 | else: |
| 64 | self.YUV = np.stack((Y, U, V), axis=-1) |
| 65 | return self.YUV, True |
| 66 | |
| 67 | def close(self): |
| 68 | self.fp.close() |
| 69 | |
| 70 |
no outgoing calls
no test coverage detected