66from torch .autograd import Variable
77
88# Set PyTorch model directory
9- os .environ ["TORCH_MODEL_ZOO " ] = "./model"
9+ os .environ ["TORCH_HOME " ] = "./model"
1010
1111squeeze = models .squeezenet1_1 (pretrained = True )
1212squeeze .eval ()
1313
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 ])
1815
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+ )
2524
26- with open (' labels.json' ) as f :
25+ with open (" labels.json" ) as f :
2726 labels_data = json .load (f )
2827
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+
3030
3131def classify_image_pytorch (image_path ):
3232
@@ -41,7 +41,6 @@ def classify_image_pytorch(image_path):
4141 for prediction in top_k :
4242 description = labels [prediction ]
4343 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 )))
4745
46+ return results
0 commit comments