danibalcells commited on
Commit
3e3cea5
1 Parent(s): 9a23f31

Remove extractor and storage attributes from IIRWI as they are already attributes of Retriever

Browse files
Files changed (1) hide show
  1. iirwi.py +4 -5
iirwi.py CHANGED
@@ -4,17 +4,16 @@ from retrieval import CroppingImgRetriever
4
  from plot import MemePlotter
5
 
6
  class IIRWI:
7
- def __init__(self, extractor, storage):
8
- self.extractor = extractor
9
- self.storage = storage
10
- self.retriever = CroppingImgRetriever(extractor, storage)
11
  self.plotter = MemePlotter()
12
 
13
  @classmethod
14
  def from_filenames(cls, extractor_filename, storage_filename):
15
  extractor = FeatureExtractor.load(extractor_filename)
16
  storage = ImgFeatureStorage.load(storage_filename)
17
- return cls(extractor, storage)
 
18
 
19
  def process(self, input_img):
20
  similar_img, crop_coords = self.retriever.process(input_img)
 
4
  from plot import MemePlotter
5
 
6
  class IIRWI:
7
+ def __init__(self, retriever):
8
+ self.retriever = retriever
 
 
9
  self.plotter = MemePlotter()
10
 
11
  @classmethod
12
  def from_filenames(cls, extractor_filename, storage_filename):
13
  extractor = FeatureExtractor.load(extractor_filename)
14
  storage = ImgFeatureStorage.load(storage_filename)
15
+ retriever = CroppingImgRetriever(extractor, storage)
16
+ return cls(retriever)
17
 
18
  def process(self, input_img):
19
  similar_img, crop_coords = self.retriever.process(input_img)