Spaces:
Sleeping
Sleeping
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}') | |