-
Notifications
You must be signed in to change notification settings - Fork 0
/
tensor3.go
151 lines (141 loc) · 2.98 KB
/
tensor3.go
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
package tensorimage
import (
"errors"
"image/color"
"gorgonia.org/tensor"
"gorgonia.org/tensor/native"
)
// dumb structure to avoid type asertion at runtime
type tensor3 struct {
c int
h int
w int
f32 [][][]float32
f64 [][][]float64
i32 [][][]int32
i64 [][][]int64
}
func toTensor3(t tensor.Tensor) (*tensor3, error) {
if len(t.Shape()) != 4 {
return nil, errors.New("TensorToImg: expected a 4D tensor (BCHW)")
}
if t.Shape()[0] != 1 {
return nil, errors.New("batch >1 not implemented")
}
dense, ok := t.(*tensor.Dense)
if !ok {
return nil, errors.New("This function can only convert dense tensors")
}
originalShape := make([]int, 4)
newShape := make([]int, 3)
copy(originalShape, t.Shape())
copy(newShape, t.Shape()[1:4])
err := dense.Reshape(newShape...)
if err != nil {
return nil, err
}
defer func() {
dense.Reshape(originalShape...)
}()
if f32, err := native.Tensor3F32(dense); err == nil {
return &tensor3{
c: dense.Shape()[0],
h: dense.Shape()[1],
w: dense.Shape()[2],
f32: f32,
}, nil
}
if f64, err := native.Tensor3F64(dense); err == nil {
return &tensor3{
c: dense.Shape()[0],
h: dense.Shape()[1],
w: dense.Shape()[2],
f64: f64,
}, nil
}
if i32, err := native.Tensor3I32(dense); err == nil {
return &tensor3{
c: dense.Shape()[0],
h: dense.Shape()[1],
w: dense.Shape()[2],
i32: i32,
}, nil
}
if i64, err := native.Tensor3I64(dense); err == nil {
return &tensor3{
c: dense.Shape()[0],
h: dense.Shape()[1],
w: dense.Shape()[2],
i64: i64,
}, nil
}
return nil, errors.New("cannot convert to tensor3")
}
func (t *tensor3) getUint8(c, h, w int) (uint8, error) {
lc := t.c
lh := t.h
lw := t.w
if c > lc || h > lh || w > lw {
return 0, errors.New("request out of bound")
}
switch {
case t.f32 != nil:
return uint8(t.f32[c][h][w]), nil
case t.f64 != nil:
return uint8(t.f64[c][h][w]), nil
case t.i32 != nil:
return uint8(t.i32[c][h][w]), nil
case t.i64 != nil:
return uint8(t.i64[c][h][w]), nil
}
return 0, nil
}
func (t *tensor3) getUint16(c, h, w int) (uint16, error) {
lc := t.c
lh := t.h
lw := t.w
if c > lc || h > lh || w > lw {
return 0, errors.New("request out of bound")
}
switch {
case t.f32 != nil:
return uint16(t.f32[c][h][w]), nil
case t.f64 != nil:
return uint16(t.f64[c][h][w]), nil
case t.i32 != nil:
return uint16(t.i32[c][h][w]), nil
case t.i64 != nil:
return uint16(t.i64[c][h][w]), nil
}
return 0, nil
}
func (t *tensor3) getColor(h, w int) (color.Color, error) {
switch t.c {
case 1:
y, err := t.getUint8(0, h, w)
return color.Gray{
Y: y,
}, err
case 3:
r, err := t.getUint8(0, h, w)
if err != nil {
return nil, err
}
g, err := t.getUint8(1, h, w)
if err != nil {
return nil, err
}
b, err := t.getUint8(2, h, w)
if err != nil {
return nil, err
}
return color.NRGBA{
R: r,
G: g,
B: b,
A: uint8(255),
}, nil
default:
return nil, errors.New("unhandled number of channel")
}
}