jiawenchim commited on
Commit
d4c4173
1 Parent(s): c6ef960

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +83 -0
app.py ADDED
@@ -0,0 +1,83 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+ model_dir = 'saved_model'
33
+ detection_model = tf.saved_model.load(str(model_dir))
34
+ return detection_model
35
+
36
+ # samples_folder = 'test_samples
37
+ # image_path = 'test_samples/sample_balloon.jpeg
38
+ #
39
+
40
+ def predict(pilimg):
41
+
42
+ image_np = pil_image_as_numpy_array(pilimg)
43
+ return predict2(image_np)
44
+
45
+ def predict2(image_np):
46
+
47
+ results = detection_model(image_np)
48
+
49
+ # different object detection models have additional results
50
+ result = {key:value.numpy() for key,value in results.items()}
51
+
52
+ label_id_offset = 0
53
+ image_np_with_detections = image_np.copy()
54
+
55
+ viz_utils.visualize_boxes_and_labels_on_image_array(
56
+ image_np_with_detections[0],
57
+ result['detection_boxes'][0],
58
+ (result['detection_classes'][0] + label_id_offset).astype(int),
59
+ result['detection_scores'][0],
60
+ category_index,
61
+ use_normalized_coordinates=True,
62
+ max_boxes_to_draw=200,
63
+ min_score_thresh=.60,
64
+ agnostic_mode=False,
65
+ line_thickness=2)
66
+
67
+ result_pil_img = tf.keras.utils.array_to_img(image_np_with_detections[0])
68
+
69
+ return result_pil_img
70
+
71
+
72
+ REPO_ID = "23A462U/iti107_cjw"
73
+ detection_model = load_model()
74
+ # pil_image = Image.open(image_path)
75
+ # image_arr = pil_image_as_numpy_array(pil_image)
76
+
77
+ # predicted_img = predict(image_arr)
78
+ # predicted_img.save('predicted.jpg')
79
+
80
+ gr.Interface(fn=predict,
81
+ inputs=gr.Image(type="pil"),
82
+ outputs=gr.Image(type="pil")
83
+ ).launch(share=True)