6
6
from torch .autograd import Variable
7
7
8
8
# Set PyTorch model directory
9
- os .environ ["TORCH_MODEL_ZOO " ] = "./model"
9
+ os .environ ["TORCH_HOME " ] = "./model"
10
10
11
11
squeeze = models .squeezenet1_1 (pretrained = True )
12
12
squeeze .eval ()
13
13
14
- normalize = transforms .Normalize (
15
- mean = [0.485 , 0.456 , 0.406 ],
16
- std = [0.229 , 0.224 , 0.225 ]
17
- )
14
+ normalize = transforms .Normalize (mean = [0.485 , 0.456 , 0.406 ], std = [0.229 , 0.224 , 0.225 ])
18
15
19
- preprocess = transforms .Compose ([
20
- transforms .Resize (256 ),
21
- transforms .CenterCrop (224 ),
22
- transforms .ToTensor (),
23
- normalize
24
- ])
16
+ preprocess = transforms .Compose (
17
+ [
18
+ transforms .Resize (256 ),
19
+ transforms .CenterCrop (224 ),
20
+ transforms .ToTensor (),
21
+ normalize ,
22
+ ]
23
+ )
25
24
26
- with open (' labels.json' ) as f :
25
+ with open (" labels.json" ) as f :
27
26
labels_data = json .load (f )
28
27
29
- labels = {int (key ):value for (key , value ) in labels_data .items ()}
28
+ labels = {int (key ): value for (key , value ) in labels_data .items ()}
29
+
30
30
31
31
def classify_image_pytorch (image_path ):
32
32
@@ -41,7 +41,6 @@ def classify_image_pytorch(image_path):
41
41
for prediction in top_k :
42
42
description = labels [prediction ]
43
43
score = fc_out .data .numpy ()[0 ][prediction ]
44
- results .append (('%s (score = %.5f)' % (description , score )))
45
-
46
- return (results )
44
+ results .append (("%s (score = %.5f)" % (description , score )))
47
45
46
+ return results
0 commit comments