Code update for inference on IT107 Assignment

app.py

@@ -7,10 +7,11 @@ from object_detection.utils import label_map_util
 from object_detection.utils import visualization_utils as viz_utils
 from object_detection.utils import ops as utils_op
 import tarfile
-import wget 
+from huggingface_hub import snapshot_download
+import os
 import gradio as gr
 
-
+MODEL_REPO = 'sokonana/it107model'
 PATH_TO_LABELS = 'data/label_map.pbtxt'   
 category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True)
 
@@ -28,10 +29,12 @@ def load_image_into_numpy_array(path):
     return pil_image_as_numpy_array(image)            
 
 def load_model():
-    wget.download("https://nyp-aicourse.s3-ap-southeast-1.amazonaws.com/pretrained-models/balloon_model.tar.gz")
-    tarfile.open("balloon_model.tar.gz").extractall()
-    model_dir = 'saved_model'    
-    detection_model = tf.saved_model.load(str(model_dir))
+    # wget.download("https://nyp-aicourse.s3-ap-southeast-1.amazonaws.com/pretrained-models/balloon_model.tar.gz")
+    # tarfile.open("balloon_model.tar.gz").extractall()
+    model_path = snapshot_download(MODEL_REPO)
+
+    model_dir = os.path.join(model_path, 'saved_model')
+    detection_model = tf.saved_model.load(model_dir)
     return detection_model    
 
 # samples_folder = 'test_samples

data/label_map.pbtxt

@@ -1,4 +1,8 @@
 item {
     id: 1
-    name: 'balloon'
+    name: 'Giant panda'
+}
+item {
+    id: 2
+    name: 'Red panda'
 }