JPniewskiHF/Object-Detection-Using-RetinaNet
0
1import gradio as gr2from huggingface_hub import from_pretrained_keras3from PIL import Image4import io5import matplotlib.pyplot as plt6import os7import re8import zipfile9import numpy as np10import tensorflow as tf11from tensorflow import keras12import tensorflow_datasets as tfds13 14coco_image = []15coco_dir = 'coco/images/'16for idx, images in enumerate(os.listdir(coco_dir)):17 image = os.path.join(coco_dir, images)18 if os.path.isfile(image) and idx < 10:19 coco_image.append(image)20 21_, dataset_info = tfds.load(22 "coco/2017", split=["train", "validation","test"], with_info=True, data_dir="data"23)24#test_dataset = tfds.load("coco/2017", split="test", data_dir="data")25int2str = dataset_info.features["objects"]["label"].int2str26 27class AnchorBox:28 """Generates anchor boxes.29 30 This class has operations to generate anchor boxes for feature maps at31 strides `[8, 16, 32, 64, 128]`. Where each anchor each box is of the32 format `[x, y, width, height]`.33 34 Attributes:35 aspect_ratios: A list of float values representing the aspect ratios of36 the anchor boxes at each location on the feature map37 scales: A list of float values representing the scale of the anchor boxes38 at each location on the feature map.39 num_anchors: The number of anchor boxes at each location on feature map40 areas: A list of float values representing the areas of the anchor41 boxes for each feature map in the feature pyramid.42 strides: A list of float value representing the strides for each feature43 map in the feature pyramid.44 """45 46 def __init__(self):47 self.aspect_ratios = [0.5, 1.0, 2.0]48 self.scales = [2 ** x for x in [0, 1 / 3, 2 / 3]]49 50 self._num_anchors = len(self.aspect_ratios) * len(self.scales)51 self._strides = [2 ** i for i in range(3, 8)]52 self._areas = [x ** 2 for x in [32.0, 64.0, 128.0, 256.0, 512.0]]53 self._anchor_dims = self._compute_dims()54 55 def _compute_dims(self):56 """Computes anchor box dimensions for all ratios and scales at all levels57 of the feature pyramid.58 """59 anchor_dims_all = []60 for area in self._areas:61 anchor_dims = []62 for ratio in self.aspect_ratios:63 anchor_height = tf.math.sqrt(area / ratio)64 anchor_width = area / anchor_height65 dims = tf.reshape(66 tf.stack([anchor_width, anchor_height], axis=-1), [1, 1, 2]67 )68 for scale in self.scales:69 anchor_dims.append(scale * dims)70 anchor_dims_all.append(tf.stack(anchor_dims, axis=-2))71 return anchor_dims_all72 73 def _get_anchors(self, feature_height, feature_width, level):74 """Generates anchor boxes for a given feature map size and level75 76 Arguments:77 feature_height: An integer representing the height of the feature map.78 feature_width: An integer representing the width of the feature map.79 level: An integer representing the level of the feature map in the80 feature pyramid.81 82 Returns:83 anchor boxes with the shape84 `(feature_height * feature_width * num_anchors, 4)`85 """86 rx = tf.range(feature_width, dtype=tf.float32) + 0.587 ry = tf.range(feature_height, dtype=tf.float32) + 0.588 centers = tf.stack(tf.meshgrid(rx, ry), axis=-1) * self._strides[level - 3]89 centers = tf.expand_dims(centers, axis=-2)90 centers = tf.tile(centers, [1, 1, self._num_anchors, 1])91 dims = tf.tile(92 self._anchor_dims[level - 3], [feature_height, feature_width, 1, 1]93 )94 anchors = tf.concat([centers, dims], axis=-1)95 return tf.reshape(96 anchors, [feature_height * feature_width * self._num_anchors, 4]97 )98 99 def get_anchors(self, image_height, image_width):100 """Generates anchor boxes for all the feature maps of the feature pyramid.101 102 Arguments:103 image_height: Height of the input image.104 image_width: Width of the input image.105 106 Returns:107 anchor boxes for all the feature maps, stacked as a single tensor108 with shape `(total_anchors, 4)`109 """110 anchors = [111 self._get_anchors(112 tf.math.ceil(image_height / 2 ** i),113 tf.math.ceil(image_width / 2 ** i),114 i,115 )116 for i in range(3, 8)117 ]118 return tf.concat(anchors, axis=0)119 120class DecodePredictions(tf.keras.layers.Layer):121 """A Keras layer that decodes predictions of the RetinaNet model.122 123 Attributes:124 num_classes: Number of classes in the dataset125 confidence_threshold: Minimum class probability, below which detections126 are pruned.127 nms_iou_threshold: IOU threshold for the NMS operation128 max_detections_per_class: Maximum number of detections to retain per129 class.130 max_detections: Maximum number of detections to retain across all131 classes.132 box_variance: The scaling factors used to scale the bounding box133 predictions.134 """135 136 def __init__(137 self,138 num_classes=80,139 confidence_threshold=0.05,140 nms_iou_threshold=0.5,141 max_detections_per_class=100,142 max_detections=100,143 box_variance=[0.1, 0.1, 0.2, 0.2],144 **kwargs145 ):146 super(DecodePredictions, self).__init__(**kwargs)147 self.num_classes = num_classes148 self.confidence_threshold = confidence_threshold149 self.nms_iou_threshold = nms_iou_threshold150 self.max_detections_per_class = max_detections_per_class151 self.max_detections = max_detections152 153 self._anchor_box = AnchorBox()154 self._box_variance = tf.convert_to_tensor(155 [0.1, 0.1, 0.2, 0.2], dtype=tf.float32156 )157 158 def _decode_box_predictions(self, anchor_boxes, box_predictions):159 boxes = box_predictions * self._box_variance160 boxes = tf.concat(161 [162 boxes[:, :, :2] * anchor_boxes[:, :, 2:] + anchor_boxes[:, :, :2],163 tf.math.exp(boxes[:, :, 2:]) * anchor_boxes[:, :, 2:],164 ],165 axis=-1,166 )167 boxes_transformed = convert_to_corners(boxes)168 return boxes_transformed169 170 def call(self, images, predictions):171 image_shape = tf.cast(tf.shape(images), dtype=tf.float32)172 anchor_boxes = self._anchor_box.get_anchors(image_shape[1], image_shape[2])173 box_predictions = predictions[:, :, :4]174 cls_predictions = tf.nn.sigmoid(predictions[:, :, 4:])175 boxes = self._decode_box_predictions(anchor_boxes[None, ...], box_predictions)176 177 return tf.image.combined_non_max_suppression(178 tf.expand_dims(boxes, axis=2),179 cls_predictions,180 self.max_detections_per_class,181 self.max_detections,182 self.nms_iou_threshold,183 self.confidence_threshold,184 clip_boxes=False,185 )186 187def convert_to_corners(boxes):188 """Changes the box format to corner coordinates189 190 Arguments:191 boxes: A tensor of rank 2 or higher with a shape of `(..., num_boxes, 4)`192 representing bounding boxes where each box is of the format193 `[x, y, width, height]`.194 195 Returns:196 converted boxes with shape same as that of boxes.197 """198 return tf.concat(199 [boxes[..., :2] - boxes[..., 2:] / 2.0, boxes[..., :2] + boxes[..., 2:] / 2.0],200 axis=-1,201 )202 203def resize_and_pad_image(204 image, min_side=800.0, max_side=1333.0, jitter=[640, 1024], stride=128.0205):206 """Resizes and pads image while preserving aspect ratio.207 208 1. Resizes images so that the shorter side is equal to `min_side`209 2. If the longer side is greater than `max_side`, then resize the image210 with longer side equal to `max_side`211 3. Pad with zeros on right and bottom to make the image shape divisible by212 `stride`213 214 Arguments:215 image: A 3-D tensor of shape `(height, width, channels)` representing an216 image.217 min_side: The shorter side of the image is resized to this value, if218 `jitter` is set to None.219 max_side: If the longer side of the image exceeds this value after220 resizing, the image is resized such that the longer side now equals to221 this value.222 jitter: A list of floats containing minimum and maximum size for scale223 jittering. If available, the shorter side of the image will be224 resized to a random value in this range.225 stride: The stride of the smallest feature map in the feature pyramid.226 Can be calculated using `image_size / feature_map_size`.227 228 Returns:229 image: Resized and padded image.230 image_shape: Shape of the image before padding.231 ratio: The scaling factor used to resize the image232 """233 image_shape = tf.cast(tf.shape(image)[:2], dtype=tf.float32)234 if jitter is not None:235 min_side = tf.random.uniform((), jitter[0], jitter[1], dtype=tf.float32)236 ratio = min_side / tf.reduce_min(image_shape)237 if ratio * tf.reduce_max(image_shape) > max_side:238 ratio = max_side / tf.reduce_max(image_shape)239 image_shape = ratio * image_shape240 image = tf.image.resize(image, tf.cast(image_shape, dtype=tf.int32))241 padded_image_shape = tf.cast(242 tf.math.ceil(image_shape / stride) * stride, dtype=tf.int32243 )244 image = tf.image.pad_to_bounding_box(245 image, 0, 0, padded_image_shape[0], padded_image_shape[1]246 )247 return image, image_shape, ratio248 249def visualize_detections(250 image, boxes, classes, scores, figsize=(7, 7), linewidth=1, color=[0, 0, 1]251):252 """Visualize Detections"""253 image = np.array(image, dtype=np.uint8)254 plt.figure(figsize=figsize)255 plt.axis("off")256 plt.imshow(image)257 ax = plt.gca()258 for box, _cls, score in zip(boxes, classes, scores):259 text = "{}: {:.2f}".format(_cls, score)260 x1, y1, x2, y2 = box261 w, h = x2 - x1, y2 - y1262 patch = plt.Rectangle(263 [x1, y1], w, h, fill=False, edgecolor=color, linewidth=linewidth264 )265 ax.add_patch(patch)266 ax.text(267 x1,268 y1,269 text,270 bbox={"facecolor": color, "alpha": 0.4},271 clip_box=ax.clipbox,272 clip_on=True,273 )274 plt.show()275 return ax276 277def prepare_image(image):278 image, _, ratio = resize_and_pad_image(image, jitter=None)279 image = tf.keras.applications.resnet.preprocess_input(image)280 return tf.expand_dims(image, axis=0), ratio281 282model = from_pretrained_keras("keras-io/Object-Detection-RetinaNet")283img_input = tf.keras.Input(shape=[None, None, 3], name="image")284predictions = model(img_input, training=False)285detections = DecodePredictions(confidence_threshold=0.5)(img_input, predictions)286inference_model = tf.keras.Model(inputs=img_input, outputs=detections)287 288def predict(image):289 input_image, ratio = prepare_image(image)290 detections = inference_model.predict(input_image)291 num_detections = detections.valid_detections[0]292 class_names = [293 int2str(int(x)) for x in detections.nmsed_classes[0][:num_detections]294 ]295 img_buf = io.BytesIO()296 ax = visualize_detections(297 image,298 detections.nmsed_boxes[0][:num_detections] / ratio,299 class_names,300 detections.nmsed_scores[0][:num_detections],301 )302 ax.figure.savefig(img_buf)303 img_buf.seek(0)304 img = Image.open(img_buf)305 return img306 307# Input308input = gr.inputs.Image(image_mode="RGB", type="numpy", label="Enter Object Image")309 310# Output311output = gr.outputs.Image(type="pil", label="Detected Objects with Class Category")312 313title = "Object Detection With RetinaNet"314description = "Upload an Image or take one from examples to localize objects present in an image, and at the same time, classify them into different categories"315 316gr.Interface(fn=predict, inputs = input, outputs = output, examples=coco_image, allow_flagging=False, analytics_enabled=False, title=title, description=description, article="<center>Space By: <u><a href='https://github.com/robotjellyzone'><b>Kavya Bisht</b></a></u> \n Based on notebook <a href='https://keras.io/examples/vision/retinanet/'><b>this notebook</b></a></center>").launch(enable_queue=True, debug=True)