Spaces:
Sleeping
Sleeping
File size: 1,401 Bytes
4094dc9 9dd1eaa 4094dc9 9dd1eaa 4094dc9 9dd1eaa 4094dc9 9dd1eaa 4094dc9 9dd1eaa 4094dc9 9dd1eaa |
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 |
import keras
import keras_cv
import numpy as np
import tensorflow as tf
from playgrounds.load_video import frames_from_video_file
def main():
pretrained_model = keras_cv.models.YOLOV8Detector.from_preset(
"yolo_v8_m_pascalvoc", bounding_box_format="xywh"
)
print('Model loaded.')
inference_resizing = keras_cv.layers.Resizing(
640, 640, pad_to_aspect_ratio=True, bounding_box_format="xywh"
)
class_ids = [
"Aeroplane", "Bicycle", "Bird", "Boat", "Bottle", "Bus", "Car", "Cat", "Chair", "Cow", "Dining Table",
"Dog", "Horse", "Motorbike", "Person", "Potted Plant", "Sheep", "Sofa", "Train", "Tvmonitor", "Total",
]
class_mapping = {i: c for (i, c) in enumerate(class_ids)}
# raw = tf.io.read_file('assets/IMG_9528.gif')
# video = tf.io.decode_gif(raw)
video = frames_from_video_file('assets/dataset/Flying/2kNjmM8BnD0_230.0_238.0.mp4', 3, (640,640))
image = video[0]
image = (image*255).astype(np.uint8)
file = tf.io.encode_png(image)
tf.io.write_file('out/t.png', file)
# image = keras.utils.load_img('assets/nick-morales-BwYcH78rcpI-unsplash.jpg')
# image = np.array(image)
image_batch = inference_resizing([image])
y_pred = pretrained_model.predict(image_batch)
classes = y_pred['classes']
boxes = y_pred["boxes"]
print(f'Classes: {classes}')
print(f'Boxes: {boxes}')
|