Update app.py

app.py

@@ -30,8 +30,8 @@ def load_image_into_numpy_array(path):
     image = Image.open(BytesIO(image_data))
     return pil_image_as_numpy_array(image)            
 
-def load_model():
-    download_dir = snapshot_download(REPO_ID)
+def load_model(model_repo_id):
+    download_dir = snapshot_download(model_repo_id)
     saved_model_dir = os.path.join(download_dir, "saved_model")
     detection_model = tf.saved_model.load(saved_model_dir)
     return detection_model
@@ -64,9 +64,40 @@ def predict2(image_np):
         agnostic_mode=False,
         line_thickness=2)
 
-    result_pil_img = tf.keras.utils.array_to_img(image_np_with_detections[0])
+    result_pil_img2 = tf.keras.utils.array_to_img(image_np_with_detections[0])
     
-    return result_pil_img
+    return result_pil_img2
+
+def predict3(pilimg):
+
+    image_np = pil_image_as_numpy_array(pilimg)
+    return predict4(image_np)
+
+def predict4(image_np):
+
+    results = detection_model2(image_np)
+
+    # different object detection models have additional results
+    result = {key:value.numpy() for key,value in results.items()}
+    
+    label_id_offset = 0
+    image_np_with_detections = image_np.copy()
+
+    viz_utils.visualize_boxes_and_labels_on_image_array(
+        image_np_with_detections[0],
+        result['detection_boxes'][0],
+        (result['detection_classes'][0] + label_id_offset).astype(int),
+        result['detection_scores'][0],
+        category_index,
+        use_normalized_coordinates=True,
+        max_boxes_to_draw=200,
+        min_score_thresh=.60,
+        agnostic_mode=False,
+        line_thickness=2)
+
+    result_pil_img4 = tf.keras.utils.array_to_img(image_np_with_detections[0])
+    
+    return result_pil_img4
 
 def detect_video(video):
     # Create a video capture object
@@ -110,7 +141,10 @@ def detect_video(video):
 
 label_id_offset = 0    
 REPO_ID = "apailang/mytfodmodel"
-detection_model = load_model()
+detection_model = load_model(REPO_ID)
+REPO_ID2 = "apailang/mytfodmodeltuned"
+detection_model2 = load_model(REPO_ID2)
+
 samples_folder = 'data'
 # pil_image = Image.open(image_path)
 # image_arr = pil_image_as_numpy_array(pil_image)
@@ -131,21 +165,32 @@ test10 = os.path.join(os.path.dirname(__file__), "data/test10.jpeg")
 test11 = os.path.join(os.path.dirname(__file__), "data/test11.jpeg")
 test12 = os.path.join(os.path.dirname(__file__), "data/test12.jpeg")
 
-tts_demo = gr.Interface(
+base_image = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
     outputs=gr.Image(type="pil"),
-    title="Image Prediction Interface",
-    description="Upload a Image for prediction",
+    title="Luffy and Chopper face detection (Base mobile net model)",
+    description="Upload a Image for prediction or click on below examples",
     examples=[[test1],[test2],[test3],[test4],[test5],[test6],[test7],[test8],[test9],[test10],[test11],[test12],],
     cache_examples=True
     )#.launch(share=True)
 
+tuned_image = gr.Interface(
+    fn=predict3,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+    title="Luffy and Chopper face detection (tuned mobile net model)",
+    description="Upload a Image for prediction or click on below examples. Mobile net tuned with data Augmentation",
+    examples=[[test1],[test2],[test3],[test4],[test5],[test6],[test7],[test8],[test9],[test10],[test11],[test12],],
+    cache_examples=True
+    )#.launch(share=True)
+
+
+
 a = os.path.join(os.path.dirname(__file__), "data/a.mp4")  # Video
 b = os.path.join(os.path.dirname(__file__), "data/b.mp4")  # Video
 c = os.path.join(os.path.dirname(__file__), "data/c.mp4")  # Video
 
-
 video_out_file = os.path.join(samples_folder,'detected' + '.mp4')
 
 stt_demo = gr.Interface(
@@ -160,7 +205,7 @@ stt_demo = gr.Interface(
     cache_examples=False
 )
 
-demo = gr.TabbedInterface([tts_demo, stt_demo], ["Image", "Video"])
+demo = gr.TabbedInterface([base_image,tuned_image, stt_demo], ["Image (base model)","Image (tuned model)", "Video"])
 
 if __name__ == "__main__":
     demo.launch()