1# -*- coding: utf-8 -*-
2
3# romiq.kh@gmail.com, 2014
4
5import array, struct, io
6
7from . import EngineError
8
9try:
10    from PIL import Image
11except ImportError:
12    Image = None
13
14class BMPLoader:
15    def __init__(self):
16        self.rgb = None
17        self.image = None
18        self.width = 0
19        self.height = 0
20
21    def load_data_int16(self, f):
22        # check magic string "BM"
23        temp = f.read(2)
24        if temp != b"BM":
25            raise EngineError("Bad magic string")
26        off = 2
27
28        temp = f.read(12)
29        f_sz, res1, res2, data_offset = struct.unpack_from("<IHHI", temp)
30        off += 12
31
32        # read next 40 bytes, BITMAPINFOHEADER
33        temp = f.read(40)
34        pict = struct.unpack_from("<IiiHHIIiiII", temp)
35        off += 40
36        if pict[0] != 40:
37            raise EngineError("Unsupported InfoHeader")
38        pictw = pict[1]
39        picth = pict[2]
40
41        # read data_offset - 40 - 6 bytes
42        delta = data_offset - 40 - 6
43        if delta < 0:
44            raise EngineError("To small bitmap data offset")
45        if delta != 8:
46            raise EngineError("Unsupported Header at 0x36")
47        temp = f.read(delta)
48        hdr36 = struct.unpack_from("<II", temp)
49        off += delta
50
51        bsz = pictw * picth * 2
52        picture_data = f.read(bsz)
53
54        off += bsz
55        if len(picture_data) != bsz:
56            raise EngineError("Bitmap truncated, need {}, got {}".format(bsz, \
57                len(picture_data)))
58
59        # read 2 zero bytes
60        #temp = f.read(2)
61        #if temp != b"\x00\x00":
62        #    raise EngineError("Magic zero bytes absent or mismatch")
63        #off += 2
64
65        temp = f.read(2)
66        if len(temp) > 0:
67            if temp != b"\x00\x00":
68                raise EngineError("BMP read error, some data unparsed")
69
70        return pictw, picth, picture_data
71
72    def pixelswap16ud(self, pw, ph, pd):
73        # convert 16 bit to 24 + vertical reverse
74        b16arr = array.array("H") # unsigned short
75        b16arr.frombytes(pd)
76        b16arr.byteswap()
77        rgb = array.array("B", [0] * pw * ph * 3)
78        for j in range(ph):
79            for i in range(pw):
80                off = (ph - j - 1) * pw * 3 + i * 3
81                b16 = b16arr[j * pw + i]
82                rgb[off] = (b16 << 3) & 0b11111000
83                rgb[off + 1] = (b16 >> 3) & 0b11111100
84                rgb[off + 2] = (b16 >> 8) & 0b11111000
85        return rgb
86
87    def pixelswap16(self, pw, ph, pd):
88        # convert 16 bit to 24
89        b16arr = array.array("H") # unsigned short
90        b16arr.frombytes(pd)
91        #b16arr.byteswap()
92        rgb = array.array("B", [0] * pw * ph * 3)
93        for j in range(ph):
94            for i in range(pw):
95                off = j * pw * 3 + i * 3
96                b16 = b16arr[j * pw + i]
97                rgb[off + 2] = (b16 << 3) & 0b11111000
98                rgb[off + 1] = (b16 >> 3) & 0b11111100
99                rgb[off + 0] = (b16 >> 8) & 0b11111000
100        return rgb
101
102    def load_info(self, f):
103        try:
104            pw, ph, pd = self.load_data_int16(f)
105            self.width = pw
106            self.height = ph
107        except:
108            f.seek(0)
109            self.image = Image.open(f)
110
111    def load_raw(self, pw, ph, pd):
112        if Image:
113            pd = self.pixelswap16(pw, ph, pd).tobytes()
114            self.image = Image.frombytes("RGB", (pw, ph), pd)
115        else:
116            self.width = pw
117            self.height = ph
118            self.rgb = self.pixelswap16(pw, ph, pd)
119
120    def load_data(self, f):
121        try:
122            pw, ph, pd = self.load_data_int16(f)
123            if Image:
124                pd = self.pixelswap16ud(pw, ph, pd).tobytes()
125                self.image = Image.frombytes("RGB", (pw, ph), pd)
126            else:
127                self.width = pw
128                self.height = ph
129                self.rgb = self.pixelswap16(pw, ph, pd)
130        except:
131            f.seek(0)
132            self.image = Image.open(f)
133