dalybuilds commited on
Commit
b04c0f2
·
verified ·
1 Parent(s): 90e019a

Update model_utils.py

Browse files
Files changed (1) hide show
  1. model_utils.py +12 -4
model_utils.py CHANGED
@@ -4,6 +4,7 @@ from transformers import ViTForImageClassification
4
  from PIL import Image
5
  import torch
6
  from dataset_utils import DatasetHandler
 
7
 
8
  class BugClassifier:
9
  def __init__(self, model_path="google/vit-base-patch16-224"):
@@ -14,9 +15,16 @@ class BugClassifier:
14
  "Japanese Beetle", "Garden Spider", "Green Grasshopper",
15
  "Luna Moth", "Common Dragonfly", "Honey Bee", "Paper Wasp"
16
  ]
17
- # Dynamically load species descriptions
18
- handler = DatasetHandler()
19
- self.species_descriptions = handler.load_descriptions(max_records=500)
 
 
 
 
 
 
 
20
 
21
  def predict(self, image):
22
  try:
@@ -39,4 +47,4 @@ class BugClassifier:
39
  def get_species_info(self, species):
40
  return self.species_descriptions.get(
41
  species, "Information not available. Consider updating your dataset for this species."
42
- )
 
4
  from PIL import Image
5
  import torch
6
  from dataset_utils import DatasetHandler
7
+ import threading
8
 
9
  class BugClassifier:
10
  def __init__(self, model_path="google/vit-base-patch16-224"):
 
15
  "Japanese Beetle", "Garden Spider", "Green Grasshopper",
16
  "Luna Moth", "Common Dragonfly", "Honey Bee", "Paper Wasp"
17
  ]
18
+ self.species_descriptions = {}
19
+ self.load_species_descriptions()
20
+
21
+ def load_species_descriptions(self):
22
+ def load():
23
+ handler = DatasetHandler()
24
+ self.species_descriptions = handler.load_descriptions(max_records=500)
25
+
26
+ thread = threading.Thread(target=load)
27
+ thread.start()
28
 
29
  def predict(self, image):
30
  try:
 
47
  def get_species_info(self, species):
48
  return self.species_descriptions.get(
49
  species, "Information not available. Consider updating your dataset for this species."
50
+ )