Spaces:
Sleeping
Sleeping
File size: 5,773 Bytes
050c87a bf269bc 050c87a c54e670 050c87a e89b7ca 050c87a a129d63 050c87a 161120c d89de93 161120c e90d845 161120c a129d63 161120c d89de93 050c87a c15f918 050c87a c77a11e 7cdb8cc c54e670 161120c c77a11e 075a84b be8fae0 c54e670 161120c f633603 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 |
import matplotlib.pyplot as plt
import numpy as np
from six import BytesIO
from PIL import Image
import tensorflow as tf
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
from tqdm import tqdm
import tarfile
import wget
import gradio as gr
from huggingface_hub import snapshot_download
import os
import cv2
PATH_TO_LABELS = 'data/label_map.pbtxt'
category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True)
label_id_offset = 0
def pil_image_as_numpy_array(pilimg):
img_array = tf.keras.utils.img_to_array(pilimg)
img_array = np.expand_dims(img_array, axis=0)
return img_array
def load_image_into_numpy_array(path):
image = None
image_data = tf.io.gfile.GFile(path, 'rb').read()
image = Image.open(BytesIO(image_data))
return pil_image_as_numpy_array(image)
def load_model():
download_dir = snapshot_download(REPO_ID)
saved_model_dir = os.path.join(download_dir, "saved_model")
detection_model = tf.saved_model.load(saved_model_dir)
return detection_model
def load_model2():
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))
return detection_model
# samples_folder = 'test_samples
# image_path = 'test_samples/sample_balloon.jpeg
#
def predict(pilimg):
image_np = pil_image_as_numpy_array(pilimg)
return predict2(image_np)
def predict2(image_np):
results = detection_model(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_img = tf.keras.utils.array_to_img(image_np_with_detections[0])
return result_pil_img
def write_video(video_in_filepath, video_out_filepath, detection_model):
video_reader = cv2.VideoCapture(video_in_filepath)
nb_frames = int(video_reader.get(cv2.CAP_PROP_FRAME_COUNT))
frame_h = int(video_reader.get(cv2.CAP_PROP_FRAME_HEIGHT))
frame_w = int(video_reader.get(cv2.CAP_PROP_FRAME_WIDTH))
fps = video_reader.get(cv2.CAP_PROP_FPS)
video_writer = cv2.VideoWriter(video_out_filepath,
cv2.VideoWriter_fourcc(*'mp4v'),
fps,
(frame_w, frame_h))
while True:
ret, image_np = video_reader.read()
if not ret:
break
results = predict(image_np)
results_np = np.array(results)
video_writer.write(results_np)
#input_tensor = tf.convert_to_tensor(np.expand_dims(image_np, 0), dtype=tf.uint8)
#results = detection_model(input_tensor)
#viz_utils.visualize_boxes_and_labels_on_image_array(
# image_np,
# results['detection_boxes'][0].numpy(),
# (results['detection_classes'][0].numpy()+ label_id_offset).astype(int),
# results['detection_scores'][0].numpy(),
# category_index,
# use_normalized_coordinates=True,
# max_boxes_to_draw=200,
# min_score_thresh=.50,
# agnostic_mode=False,
# line_thickness=2)
#video_writer.write(np.uint8(image_np))
# Release camera and close windows
video_reader.release()
video_writer.release()
cv2.destroyAllWindows()
cv2.waitKey(1)
def predict_video (video_file_name):
detected_video_file = "detected_video.mp4"
write_video(video_file_name,detected_video_file,detection_model)
return detected_video_file
REPO_ID = "YEHTUT/tfodmodel"
detection_model = load_model()
# pil_image = Image.open(image_path)
# image_arr = pil_image_as_numpy_array(pil_image)
# predicted_img = predict(image_arr)
# predicted_img.save('predicted.jpg')
Image_tab = gr.Interface(fn=predict,
inputs=gr.Image(type="pil"),
outputs=gr.Image(type="pil"),
examples=[["SampleImage1.jpg"],["SampleImage2.jpg"],["SampleImage3.jpg"],["SampleImage4.jpg"],["SampleImage5.jpg"],["SampleImage6.jpg"]],
title="This is the object detection model for Durian and Pineapple images",
description="Using ssd_mobilenet_v2_320x320_coco17_tpu-8 to detect Durina and Pineapple"
)
Video_tab = gr.Interface(fn=predict_video,
inputs=gr.Video(label="Upload Video"),
outputs=gr.Video(label="Detected Video"),
examples=[["SampleVideo1.mp4"]],
title="This is the object detection model for Durian and Pineapple videos",
description="Using ssd_mobilenet_v2_320x320_coco17_tpu-8 to detect Durina and Pineapple"
)
gr.TabbedInterface([Image_tab, Video_tab], ["Image", "Video"]).launch(share=True)
#gr.Interface(fn=predict,
# inputs=gr.Image(type="pil"),
# outputs=gr.Image(type="pil")
# ).launch(share=True)
|