ksh-nyp commited on
Commit
8c5c96d
1 Parent(s): 5dcd0de

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +80 -0
app.py ADDED
@@ -0,0 +1,80 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import matplotlib.pyplot as plt
2
+ import numpy as np
3
+ from six import BytesIO
4
+ from PIL import Image
5
+ import tensorflow as tf
6
+ from object_detection.utils import label_map_util
7
+ from object_detection.utils import visualization_utils as viz_utils
8
+ from object_detection.utils import ops as utils_op
9
+ import tarfile
10
+ import wget
11
+ import gradio as gr
12
+ from huggingface_hub import snapshot_download
13
+ import os
14
+
15
+ PATH_TO_LABELS = 'data/label_map.pbtxt'
16
+ category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True)
17
+
18
+ def pil_image_as_numpy_array(pilimg):
19
+
20
+ img_array = tf.keras.utils.img_to_array(pilimg)
21
+ img_array = np.expand_dims(img_array, axis=0)
22
+ return img_array
23
+
24
+ def load_image_into_numpy_array(path):
25
+
26
+ image = None
27
+ image_data = tf.io.gfile.GFile(path, 'rb').read()
28
+ image = Image.open(BytesIO(image_data))
29
+ return pil_image_as_numpy_array(image)
30
+
31
+ def load_model():
32
+ download_dir = snapshot_download(REPO_ID)
33
+ saved_model_dir = os.path.join(download_dir, "saved_model")
34
+ detection_model = tf.saved_model.load(saved_model_dir)
35
+ return detection_model
36
+
37
+ def predict(pilimg):
38
+
39
+ image_np = pil_image_as_numpy_array(pilimg)
40
+ return predict2(image_np)
41
+
42
+ def predict2(image_np):
43
+
44
+ results = detection_model(image_np)
45
+
46
+ # different object detection models have additional results
47
+ result = {key:value.numpy() for key,value in results.items()}
48
+
49
+ label_id_offset = 0
50
+ image_np_with_detections = image_np.copy()
51
+
52
+ viz_utils.visualize_boxes_and_labels_on_image_array(
53
+ image_np_with_detections[0],
54
+ result['detection_boxes'][0],
55
+ (result['detection_classes'][0] + label_id_offset).astype(int),
56
+ result['detection_scores'][0],
57
+ category_index,
58
+ use_normalized_coordinates=True,
59
+ max_boxes_to_draw=200,
60
+ min_score_thresh=.60,
61
+ agnostic_mode=False,
62
+ line_thickness=2)
63
+
64
+ result_pil_img = tf.keras.utils.array_to_img(image_np_with_detections[0])
65
+
66
+ return result_pil_img
67
+
68
+
69
+ REPO_ID = "ksh-nyp/TFODModel-1-LeoD"
70
+ detection_model = load_model()
71
+ # pil_image = Image.open(image_path)
72
+ # image_arr = pil_image_as_numpy_array(pil_image)
73
+
74
+ # predicted_img = predict(image_arr)
75
+ # predicted_img.save('predicted.jpg')
76
+
77
+ gr.Interface(fn=predict,
78
+ inputs=gr.Image(type="pil"),
79
+ outputs=gr.Image(type="pil")
80
+ ).launch(share=True)