<s> import os <EOL> import tarfile <EOL> import cv2 <EOL> import numpy as np <EOL> from ... utils import logger , get_rng <EOL> from . . base import DataFlow <EOL> from ... utils . fs import mkdir_p , download <EOL> __all__ = [ '<STR_LIT>' , '<STR_LIT>' ] <EOL> CAFFE_ILSVRC12_URL = "<STR_LIT>" <EOL> CAFFE_PROTO_URL = "<STR_LIT>" <EOL> """<STR_LIT>""" <EOL> class ILSVRCMeta ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , dir = None ) : <EOL> if dir is None : <EOL> dir = os . path . join ( os . path . dirname ( __file__ ) , '<STR_LIT>' ) <EOL> self . dir = dir <EOL> mkdir_p ( self . dir ) <EOL> self . caffe_pb_file = os . path . join ( self . dir , '<STR_LIT>' ) <EOL> if not os . path . isfile ( self . caffe_pb_file ) : <EOL> self . _download_caffe_meta ( ) <EOL> def get_synset_words_1000 ( self ) : <EOL> """<STR_LIT>""" <EOL> fname = os . path . join ( self . dir , '<STR_LIT>' ) <EOL> assert os . path . isfile ( fname ) <EOL> lines = [ x . strip ( ) for x in open ( fname ) . readlines ( ) ] <EOL> return dict ( enumerate ( lines ) ) <EOL> def _download_caffe_meta ( self ) : <EOL> fpath = download ( CAFFE_ILSVRC12_URL , self . dir ) <EOL> tarfile . open ( fpath , '<STR_LIT>' ) . extractall ( self . dir ) <EOL> proto_path = download ( CAFFE_PROTO_URL , self . dir ) <EOL> ret = os . system ( '<STR_LIT>' . format ( self . dir ) ) <EOL> assert ret == <NUM_LIT:0> , "<STR_LIT>" <EOL> def get_image_list ( self , name ) : <EOL> """<STR_LIT>""" <EOL> assert name in [ '<STR_LIT:train>' , '<STR_LIT>' , '<STR_LIT:test>' ] <EOL> fname = os . path . join ( self . dir , name + '<STR_LIT>' ) <EOL> assert os . path . isfile ( fname ) <EOL> with open ( fname ) as f : <EOL> ret = [ ] <EOL> for line in f . readlines ( ) : <EOL> name , cls = line . strip ( ) . split ( ) <EOL> ret . append ( ( name , int ( cls ) ) ) <EOL> return ret <EOL> def get_per_pixel_mean ( self , size = None ) : <EOL> """<STR_LIT>""" <EOL> import imp <EOL> caffepb = imp . load_source ( '<STR_LIT>' , self . caffe_pb_file ) <EOL> obj = caffepb . BlobProto ( ) <EOL> mean_file = os . path . join ( self . dir , '<STR_LIT>' ) <EOL> with open ( mean_file ) as f : <EOL> obj . ParseFromString ( f . read ( ) ) <EOL> arr = np . array ( obj . data ) . reshape ( ( <NUM_LIT:3> , <NUM_LIT> , <NUM_LIT> ) ) . astype ( '<STR_LIT>' ) <EOL> arr = np . transpose ( arr , [ <NUM_LIT:1> , <NUM_LIT:2> , <NUM_LIT:0> ] ) <EOL> if size is not None : <EOL> arr = cv2 . resize ( arr , size [ : : - <NUM_LIT:1> ] ) <EOL> return arr <EOL> class ILSVRC12 ( DataFlow ) : <EOL> def __init__ ( self , dir , name , meta_dir = None , shuffle = True ) : <EOL> """<STR_LIT>""" <EOL> assert name in [ '<STR_LIT:train>' , '<STR_LIT:test>' , '<STR_LIT>' ] <EOL> self . dir = dir <EOL> self . name = name <EOL> self . shuffle = shuffle <EOL> self . meta = ILSVRCMeta ( meta_dir ) <EOL> self . imglist = self . meta . get_image_list ( name ) <EOL> self . rng = get_rng ( self ) <EOL> def size ( self ) : <EOL> return len ( self . imglist ) <EOL> def reset_state ( self ) : <EOL> """<STR_LIT>""" <EOL> self . rng = get_rng ( self ) <EOL> def get_data ( self ) : <EOL> """<STR_LIT>""" <EOL> idxs = np . arange ( len ( self . imglist ) ) <EOL> if self . shuffle : <EOL> self . rng . shuffle ( idxs ) <EOL> for k in idxs : <EOL> tp = self . imglist [ k ] <EOL> fname = os . path . join ( self . dir , self . name , tp [ <NUM_LIT:0> ] ) . strip ( ) <EOL> im = cv2 . imread ( fname ) <EOL> assert im is not None , fname <EOL> if im . ndim == <NUM_LIT:2> : <EOL> im = np . expand_dims ( im , <NUM_LIT:2> ) . repeat ( <NUM_LIT:3> , <NUM_LIT:2> ) <EOL> yield [ im , tp [ <NUM_LIT:1> ] ] <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> meta = ILSVRCMeta ( ) <EOL> print ( meta . get_per_pixel_mean ( ) ) </s>
