$ r5 K( a$ v' |; Z# |1 d img_name = self.imgs[idx]" F/ s( V' b) O
image_path = os.path.join(self.image_dir, img_name) z: s" W4 X5 i3 A) W8 E" H
7 l6 {: f: A; j! |( k; @# Z7 P7 H/ Q8 W& W; o
# reading the images and converting them to correct size and color W* p3 Z; ^6 K6 W' M- P$ O# h img = cv2.imread(image_path) 4 M/ ?. d2 n0 N5 A8 z: j l6 N img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32)% C2 y3 m9 w! U5 j. W
img_res = cv2.resize(img_rgb, (self.width, self.height), cv2.INTER_AREA) + o4 c% C: c# Q9 Q # diving by 2555 N0 Y9 Q# ~3 G9 G1 g+ R* ~( `
img_res /= 255.0 ) c3 [1 j3 ~/ b, c& @ : ] ?: i; h5 p# W$ `: M c% z , ~, O) @# O1 c # annotation file * Q! b* R5 |3 {! R3 p annot_filename = img_name[:-4] + '.xml'7 }' ~# ]/ p& H3 B. Q7 R5 N3 F
annot_file_path = os.path.join(self.xml_dir, annot_filename) / J9 R1 p, W8 F+ w. } ! d' r# S+ L8 ] $ i( y) r% G/ c* H& } ^ boxes = []7 G+ d' Q* B* R3 B0 X& K
labels = [] & K) m5 r0 P* [/ |, B% n; m$ h tree = et.parse(annot_file_path)+ a: s' n8 A" W- `# ?6 y, P
root = tree.getroot()3 g }- G# T5 Y4 A, ?
( @, ]) g1 R' } D 9 X! l8 @. ]+ z1 D* v* K8 K # cv2 image gives size as height x width6 U' d, Y8 D6 Z. c* g( E
wt = img.shape[1]6 n4 [* m6 P. H; f5 v6 U8 f
ht = img.shape[0]7 ^1 I9 @8 [ `8 o6 c
: }0 i$ E3 m8 M+ Y. G8 W , y! F7 ~' I) _2 `4 Q- N7 g # box coordinates for xml files are extracted and corrected for image size given , f8 f3 W2 O; i3 k5 X; V) L* j. s6 e for member in root.findall('object'): $ k" l1 c$ H) u7 m$ I+ W labels.append(self.classes.index(member.find('name').text))) J# C0 y& J/ ] O' Y8 }* z/ Y8 ~
, I" U6 H# h& r ; V8 g/ x; Z+ n" ]$ P # bounding box ) T% `4 ~$ U$ k& }7 D$ q, S' Y z xmin = int(member.find('bndbox').find('xmin').text) $ f) c6 j! a9 D9 c xmax = int(member.find('bndbox').find('xmax').text)" |/ R7 h0 P a* _/ t. u
2 P) s' n; Q/ A/ E& ~
% y& m' W9 J/ u/ R5 ? ymin = int(member.find('bndbox').find('ymin').text)+ [' D: F+ r4 S G; V
ymax = int(member.find('bndbox').find('ymax').text) 0 S% i9 r6 ~& f: s0 p" H . j/ [1 c# d7 T7 [) i$ |5 I ; G* ?8 I! M6 B# Q* I% C xmin_corr = (xmin / wt) * self.width % w8 i9 Y# P' h. L1 n; ~- a) v xmax_corr = (xmax / wt) * self.width $ [4 w" R; r/ H0 t# E ymin_corr = (ymin / ht) * self.height9 O. b$ P! ^1 J
ymax_corr = (ymax / ht) * self.height& h" p% I1 D4 s8 A4 t
boxes.append([xmin_corr, ymin_corr, xmax_corr, ymax_corr]) & a5 v- g$ u# H1 t' O( W8 i3 F$ p6 e, k; Y9 X
7 A% |1 X2 W& q' J- S$ n* ~* B
# convert boxes into a torch.Tensor 1 v8 b2 |8 Y% ^ boxes = torch.as_tensor(boxes, dtype=torch.float32)3 l3 @- q$ t9 k" K
. L8 D8 y ?7 y- a9 g1 N, r
! P! r2 v8 w0 L% d8 N
# getting the areas of the boxes5 Y* e; ~! k' _7 }
area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])' S$ S3 r# r9 ^2 C$ }( S
/ p6 v/ N/ G+ z3 ?- J' O7 y& a, ? `% r# |9 N% Z# W
# suppose all instances are not crowd' v7 L* e, Z7 X9 q. T
iscrowd = torch.zeros((boxes.shape[0],), dtype=torch.int64) 1 x/ G" a V! \ # \ j$ m) _, [4 O# _3 Z8 X3 X2 F* C2 e2 ^8 \& v
labels = torch.as_tensor(labels, dtype=torch.int64) ; o& V9 `3 f" s, d2 h( o) K: {4 x9 e N3 R) _* ~2 X! R% c9 A* `% }
0 a* A4 F% k1 M( P6 \+ \7 J target = {}1 y6 M5 I. l& H- W# d: u0 x
target["boxes"] = boxes W4 w, f/ W6 X, F target["labels"] = labels ( Y+ U, ]- N3 w. Q target["area"] = area ! M$ W, [4 V" l! K1 F target["iscrowd"] = iscrowd 6 g3 T. P _( \+ ?) d& G& X- A, m # image_id# l5 t- O& K6 A% k: G
image_id = torch.tensor([idx])0 V' |3 a" {6 e1 |# Q) X
target["image_id"] = image_id( d; p3 \3 |5 j; O/ B& B; T2 b9 F
$ y8 j* H. `/ P& v7 H' |
) W/ V1 S: a- ]7 r
if self.transforms: . `. \3 H8 s1 C, N8 W sample = self.transforms(image=img_res,. J1 j- `' p. y3 p6 X
bboxes=target['boxes'], ) b: t# i6 B5 j labels=labels)* _8 y6 ^' M8 {" H' f' T0 I
: e G% Q, a5 A, q. v+ Y$ [
& j t# P: Q' p; H
img_res = sample['image'] ) C- G7 D8 O5 @2 I6 _ target['boxes'] = torch.Tensor(sample['bboxes']) 7 a# T6 O! Y) K 8 T" I! [ t6 @5 k; n' _4 e $ r2 W+ P1 P5 ^ return img_res, target- E) U- T+ u& s! x4 O
. ]/ J( Q: @" o
' q" p; J4 }% z1 ^( h$ {
def __len__(self):( R n' N b1 [: R2 K% p6 p
return len(self.imgs)2 |4 y/ s( o: ?% e8 `
( h. O3 S# d: `, v
3 e. G4 I# a0 O7 Q3 Y3 H/ v
# function to convert a torchtensor back to PIL image " E) f/ E7 k! ^. s/ \+ I! }def torch_to_pil(img):9 V+ ~3 M5 s4 F8 u+ Q$ k
return torchtrans.ToPILImage()(img).convert('RGB') 6 `. |, i4 M# g; f5 b1 T8 U7 f! Y9 |6 K# [$ c% t
. @: X$ K5 a# d- t
5 \8 K) N- p2 R 8 X! O' a! B: k8 X8 Hdef plot_img_bbox(img, target): + v x0 r0 _: l- j- f9 G2 K # plot the image and bboxes. h- f: |; c, c' L
fig, a = plt.subplots(1, 1) - q% Q e0 s) t6 ^; c2 k! t/ i1 N9 P fig.set_size_inches(5, 5): }/ w5 q" v2 V! F
a.imshow(img) 3 P( K7 y& r2 }$ J7 r" ?) W: J' h for box in (target['boxes']): 1 I6 R0 S; |9 T9 Q x, y, width, height = box[0], box[1], box[2] - box[0], box[3] - box[1]* h" H- S" d: x- Y
rect = patches.Rectangle((x, y), 2 Q* h, p- k" q# D width, height, 5 C: z* l; u: M" V, F. o linewidth=2, 5 D- ^8 L' t2 C' p edgecolor='r', ; q% _3 R b4 A" B facecolor='none')( p2 A1 N2 [3 W3 o; E
: z2 l+ i3 `0 k! p ' X3 i7 [& I6 B # Draw the bounding box on top of the image N( ^. v/ r8 U, b% N3 B, \
a.add_patch(rect) ) R' K( U! ~, K% }* g4 a plt.show()" m7 g( [: h# e
1 K6 a a# t- \7 s, S3 l; b( ~
* W/ e7 m( }* {" O9 v' t : O, X# B( J9 }% g! F : m1 a( I# X( [* [def get_transform(train): , T$ H9 {) u! L6 R if train:: m- w7 S, b2 D0 ^
return A.Compose([# o; f# |& k' [: F( G# W) `6 H, W
A.HorizontalFlip(0.5), # t1 P% x8 q# ]& e) N # ToTensorV2 converts image to pytorch tensor without div by 2553 I7 g- H$ G. J+ n
ToTensorV2(p=1.0) % P; @- P1 B, O9 V5 L ], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']}) + U0 z' n7 K$ p else: % \3 }+ ~" [* v return A.Compose([ 9 ?" g9 w! P7 h8 C ToTensorV2(p=1.0), K, I4 {- T5 F3 B9 t
], bbox_params={'format': 'pascal_voc', 'label_fields': ['labels']})! b; Y4 G! B+ }* v) n
0 Y+ C0 S% V; Y1 I* o' m, i% w
" ?. |. y, F: o
" |8 O5 v4 d6 r- Q/ Y
5 _1 I7 I0 d% `" x( H, D. {8 I1 m
2 K9 N4 a6 R( Y# Z& j" J- n4 N/ S# F7 W$ Z
dataset = FruitImagesDataset(train_image_dir,train_xml_dir, 480, 480, transforms= get_transform(train=True)) " m! b; c% k. Q5 x* U" z% r3 ], ^, A! A3 G0 T' a) ~3 a
, i Y$ ]5 ~ @# Z
print(len(dataset)) k: t+ i2 a! A7 N0 k3 l4 }# getting the image and target for a test index. Feel free to change the index.3 Z' ^& p+ @4 _* r9 m
img, target = dataset[29]& E1 f( o& \; o& x9 N0 n( ~/ D
print(img.shape, '\n', target) * I, f& U! y+ bplot_img_bbox(torch_to_pil(img), target) C9 t; n1 ^. E1 # X6 _% h: u+ f8 K# ?22 L9 z/ N$ S" j) q2 Q
35 {- A* \; K" A/ n) a8 S
4: o( [% B. v5 j1 S& v; r# Z
5$ O2 ?, x/ K- ~$ k8 Y- S
6 t& Z* A& U6 I; v
7 8 O9 `! W! |0 Q, @* A81 A6 h, ~% h2 R; R9 M t
9 , ? k0 D- `) i: l- Q: m' Z N10. h4 E( \' I7 Q2 L# z
11 h) E' e8 C# J- z+ a5 w" j& U12. O8 g9 x! n$ `0 b
13 7 g$ O% y: a( f9 D& Z, f143 W" Z9 Z6 @& {5 c0 o4 e! Q J- [
15 " y* X6 t3 N, P" k1 @16 , p- ~4 q" S1 U) o1 A17 6 a7 N" J: _' M+ l189 ~6 a3 ?" {/ C: N7 O; x% t6 Q
196 f) D# ~- }3 l0 L
202 x* h: A+ F9 D- Y8 ]& s9 W1 _
21 1 \1 o- y1 [" Z7 ]$ V& G227 I$ h8 s# n( @5 E4 ?& | P
23; x7 ]* b: ^: o$ d* O* i1 Q% P
24 7 u+ e7 e6 N1 F# Z9 v* }( @* {% J25 / h; Y/ D" Z2 s' m26 0 m: O1 ?# v& y( v: t27 7 e( H- T# ]/ P- c5 e! {7 x28& }! W) W) q: n$ R
29; I- b; \4 x# x2 H Z- r
30 1 T7 d7 H. D7 J W+ y31 ) Z3 l+ E" q4 B& R; o) V: c321 s2 y& h% ~. Z* @4 j6 g
33* L6 O3 y% |2 g9 y: J
34+ x7 ^* g0 v; P! f y
35 N( M% f5 D* X) q# L1 }
360 n' O8 A1 Q$ U+ }$ @) k1 J+ p1 k
37 ! Z5 F1 q- t- N: f38 / M7 i; L" B4 r. O* q39 3 E0 E- w3 G: F# j7 Y P40 : R: @/ }# P. m& d& G, ^+ H. q. P416 }, |7 k4 ?5 ]/ L4 I _7 J2 t
423 W1 w" N$ n! o
43 7 V# ?: _7 p7 a2 o44 , x4 p4 L$ Y' E2 c4 x) }45 1 e8 t) e- u4 u46 7 _# r0 q5 W, F5 B47 2 v) }+ v8 B% a4 I+ R48 4 L7 [6 j* k5 j& }" ]49 6 \/ ~" L% r/ O0 a. e' s. `* f50 - I- Y% ]8 a$ J$ Q! K# |1 [510 g3 [! k5 E) n0 _
52+ S/ M7 ]6 z9 c2 H# \ E6 L! b
53 " Q( h; R/ c+ f O5 ?9 a54% k% B8 S/ r; s
55& y; S; C9 \5 ?" Y6 C
56, u l( Y* B# x$ d! z8 z; j
57 4 _6 Y6 J4 K/ W4 l: W588 z3 L# p X! v/ j3 b0 T& B
595 J+ v7 u/ k- {* M: D) M
60 * H* w1 C2 \3 Q M! E$ O61 8 N4 q+ D. O3 o. I: z+ `62 * T! O. y0 w+ e7 i! p0 R63 $ A9 |0 |7 R: c: @7 l64 7 @1 j) e! u( t2 Z! }65 0 z" O* \+ t* G( \6 B, y; O66 3 N# j+ C; n/ y/ r# t/ Y4 D67 $ A0 M+ f/ K" P& a) Y* B689 k1 K5 `, O0 n
691 B9 M9 p& X3 s7 h' Q1 r
70 + u) W/ k" N0 R+ G: p71 1 h5 y1 L: \7 z, q72 * C5 c( I. E: b& o u9 M6 }73 9 `% a9 y) Z$ L74. {, b9 a$ T6 m# x% }" Q R [1 T
75 " d% S8 A8 V0 r8 b) A% X9 t76 , ]2 Z& p4 B, r( s77) E1 o2 f- M6 u0 |1 i, a6 x! r
78) M9 q" p' T/ _9 }
79 ' }1 g1 B* b; Q( \6 x800 u$ C' a, {0 k& t# P
81. G! Z" b; @% Y: P5 N9 G
82 , \2 g& J& h* |2 F& I7 L833 P; v- ]$ ]0 A' k) D1 ~/ k
842 \8 |! L |1 s) M; |
85 : Q" K/ S2 A* X `: J- _86 3 B# E% l; N5 q6 \' D5 O' O2 u* _876 X2 ~ B( ? u5 x9 s
88" W; P4 R9 i: Q* z, W% K( `* G
89 3 O0 ]3 n1 {. Z; V; \: G/ R904 W ?5 e8 ~4 F N, Q8 ]3 C, n! H
910 T& H: v, w/ n6 ?$ c' T$ |4 h
92: R1 Q4 ~3 ^9 r: R+ o
937 z$ i- w6 W8 H3 {( T+ @" j# h
94 9 V: C4 i. m% r: \7 d$ W8 M95 $ E) C& v) d* w5 ?0 R96 R( h9 N7 |8 }% X& p' j0 N/ s
971 W/ ` @5 m4 p" t6 ]+ }5 h
98 6 ~2 f3 v2 f' R- U99 # \8 t7 ?$ x# Q1 q H100 - {3 q! T+ `3 X! R: P7 ?1013 T Z- D- `2 O0 M" |4 t
102 ( G1 \& u: f7 ~1033 l/ Y. D7 W$ k# c& s- A: N
1043 g7 i3 y6 ?0 U" v2 x/ ~/ Q$ j
105 % P7 R) w! U$ K: Y* F% K5 }106 6 i+ A+ j# O* ^: m7 I7 B1070 I' R/ e) V- T- W9 y
108 $ w% H" `9 ]7 l. e4 i2 i" }6 T+ G7 H7 i# N109. n7 k, M% {" V
1104 q, N- x! U1 n( ]8 v2 [
111 ( ^/ u% g/ U4 x- i1127 k4 O6 |+ I) L" ]4 r/ ]# C
113& S/ G. z5 O# r: U$ k$ q$ Y$ \
1148 p7 W0 M3 k$ ]7 Z$ n* b
115 ) m5 k- P7 B2 `6 ~4 o6 ~* |116 F9 |$ P& a) a117 . A* V( D* W% l0 \; k `9 x7 r118+ ^ s' @7 F2 d( g
119 . Z# ?3 g6 }* E# w120! ?, d. {( E% b ~+ o; k2 ~3 b: Y
121. ~# I& H# g0 v n
122 ; A6 e0 f: v9 y7 I1 V123* d8 O" \9 m9 k% d \; }
124 ; X2 Q% K8 L8 p7 l125# X9 i4 @) L0 c
126 + x1 A, v/ B; _9 r: Z127- Q& n! _2 |! W4 Y7 F. d8 v; J
128 : a1 f) j: y; ]6 i* o129 ; b4 ^# X$ W/ Z6 U5 ]( ?0 `130 5 o3 B/ Z% ~% F9 x1310 r0 ~7 p3 w7 f) ~" Z
132 8 h4 k- ^2 ?6 r+ }4 N, Y1 O4 }133 ( n/ j. C( |# h5 \) M; x134 & j9 T. ?9 d. j1359 }$ }; ~" W$ p% J
136 0 R$ c s( `# F; `# J& w' W* o137# R. P( X; W+ a- A! r7 Z; j7 u
138 ( B. j" N! n, c7 N/ \ |139 `) h2 z7 t& |7 ~9 O* v, V1402 i2 U. S7 l$ v0 \9 ]6 Y
141 7 z& ^2 y/ X. T142' X- |' e' e: a- y1 x2 D
143 ' F$ e/ \# h5 `$ Y144" Y, L; J$ E, d, g/ e H
145! ?* [0 ?% q7 o9 U# l$ R
146, r9 L: {2 F! ^; A- @+ T9 V
147 9 N+ u6 U9 n" V% d8 C% s+ _' O148+ V/ e$ ?7 [. j2 h" @, J
149/ ~9 g- F4 |- |" q; u O- w
1503 ^; q; k) e# x
151 6 @# n2 o8 F! S! S% R9 L& Q152 1 r/ q9 k2 I$ E1 c& l2 I1 q5 | E153 @5 c- B: B% P, B. `& c154 $ @: E3 V' d+ q# L155 ! y( @. A, ]6 q r/ j156 ; K7 ^1 ^8 x) R& D z输出如下: 3 {9 h3 i7 V2 I( s# o/ j( l: U- R : ]' S; Q$ h0 I; y3 \ 5 o. d; N% s* v% p6 ftorch.Size([3, 480, 480]) 9 j5 s* ]8 I, z$ J0 X
{'boxes': tensor([[130.8000, 97.8000, 327.6000, 292.2000],6 C2 P% x5 \0 T# ?) j
[159.0000, 268.8000, 349.8000, 427.8000],; Z0 G, b' {& l) [8 U
[ 0.0000, 282.0000, 118.2000, 429.6000], " S6 i M% R, T7 O [ 43.8000, 107.4000, 199.2000, 280.2000], 5 u$ H$ Y) T. p2 Z [295.2000, 37.8000, 479.4000, 248.4000]]), 'labels': tensor([0, 0, 0, 0, 0]), 'area': tensor([38257.9258, 30337.2012, 17446.3223, 26853.1270, 38792.5195]), 'iscrowd': tensor([0, 0, 0, 0, 0]), 'image_id': tensor([29])}/ |* @) x0 P( M6 k; k3 u% H8 V
1+ J, ~5 \# t$ k5 n' W; a! d4 e
20 ~) ~( g" B$ q0 i! x( M0 a5 l0 E
32 j! F+ O- Q- m! O4 b' s6 t& ]
4 " m( x( T2 l& B& ]' l. j" z+ p' G5 # r8 S, o1 E" D2 u: P6 - [) B% x; S* `+ ^) q& H1 P8 H1 G6 z% v- C3 O! R4 |