farkmu45 commited on
Commit
e4f4aa4
1 Parent(s): 4119868

Only predict one image at a time

Browse files
Files changed (3) hide show
  1. .gitignore +2 -0
  2. Processor.py +2 -12
  3. app.py +4 -1
.gitignore ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ *.pyc
2
+ yolov5x6.pt
Processor.py CHANGED
@@ -10,18 +10,8 @@ class Processor():
10
  self.__model = torch.hub.load(
11
  'ultralytics/yolov5', 'yolov5x6', trust_repo=True)
12
 
13
- def classify_images(self, images: List[NDArray]) -> List[str]:
14
- result = []
15
- class_names = self.__inference.dls.vocab
16
-
17
- test_dl = self.__inference.dls.test_dl(images)
18
- tensors = self.__inference.get_preds(dl=test_dl, with_decoded=True)[2]
19
- preds = [int(t.item()) for t in tensors]
20
-
21
- for i in preds:
22
- result.append(class_names[i])
23
-
24
- return result
25
 
26
  def filter_image(self, image: Image) -> bool:
27
  results = self.__model(image)
 
10
  self.__model = torch.hub.load(
11
  'ultralytics/yolov5', 'yolov5x6', trust_repo=True)
12
 
13
+ def classify_image(self, images: NDArray) -> str:
14
+ return self.__inference.predict(images)
 
 
 
 
 
 
 
 
 
 
15
 
16
  def filter_image(self, image: Image) -> bool:
17
  results = self.__model(image)
app.py CHANGED
@@ -14,13 +14,16 @@ def initialize_app():
14
 
15
  def process_images(images, processor: Processor):
16
  filtered_images = []
 
17
 
18
  for image in images:
19
  image = Image.open(image)
20
  if processor.filter_image(image):
21
  filtered_images.append(np.asarray(image))
 
 
 
22
 
23
- result = processor.classify_images(filtered_images)
24
  outfit = mode(result)
25
 
26
  with open(f'./texts/{outfit}.txt') as text:
 
14
 
15
  def process_images(images, processor: Processor):
16
  filtered_images = []
17
+ result = []
18
 
19
  for image in images:
20
  image = Image.open(image)
21
  if processor.filter_image(image):
22
  filtered_images.append(np.asarray(image))
23
+
24
+ for img in filtered_images:
25
+ result.append(processor.classify_image(img)[0])
26
 
 
27
  outfit = mode(result)
28
 
29
  with open(f'./texts/{outfit}.txt') as text: