gu5ousa/mars-soil-sample-localization

1import tensorflow as tf2from tensorflow import keras3from keras.layers import Conv2D, MaxPooling2D, Dense, Flatten, Dropout4from keras.models import Sequential5import cv26import numpy as np7import gradio as gr8import copy9 10 11model_regression = Sequential([12 13    Conv2D(filters=16,kernel_size=(3,3), input_shape = (384, 512, 1), activation='relu'),14    MaxPooling2D(pool_size=(2,2)),15 16    Conv2D(filters=32, kernel_size=(3,3), activation='relu'),17    MaxPooling2D(pool_size=(2,2)),18    Dropout(0.25),19 20    Conv2D(filters=64, kernel_size=(3,3), activation='relu'),21    MaxPooling2D(pool_size=(2,2)),22    Dropout(0.25),23 24    Flatten(),25    Dense(128, activation='relu'),26    Dropout(0.25),27 28    Dense(4, activation='sigmoid')29])30 31 32 33model_detection = Sequential([34 35    Conv2D(filters=64,kernel_size=(3,3),  input_shape = (360, 512, 1),activation='relu'),36    MaxPooling2D(pool_size=(2,2),strides=(2, 2),padding="valid"),37    Dropout(0.1),38 39    Conv2D(filters=32,kernel_size=(3,3), activation='relu'),40    MaxPooling2D(pool_size=(2,2),strides=(2, 2),padding="valid"),41    Dropout(0.1),42 43    Conv2D(filters=16,kernel_size=(3,3), activation='relu'),44    MaxPooling2D(pool_size=(2,2),strides=(2, 2),padding="valid"),45    Dropout(0.1),46 47    Flatten(),48    Dense(128, activation='relu'),49    Dense(2, activation='softmax')50])51 52model_regression.load_weights('grayCkpt1_17_0.00889_.h5')                                       53model_detection.load_weights('grayCkpt1_13_0.68606_.h5')54 55 56 57def show_image_bbox(X_arr, y_arr, confidence, new_img_path):58 59  y_arr_albu = [0,0,0,0]60 61  y_arr_albu[0] = y_arr[0] - y_arr[2] / 262  y_arr_albu[1] = y_arr[1] - y_arr[3] / 263  y_arr_albu[2] = y_arr[2] + y_arr_albu[0]64  y_arr_albu[3] = y_arr[3] + y_arr_albu[1]65 66 67  shape_after = X_arr.shape68 69  X_arr = cv2.merge([X_arr, X_arr, X_arr])   70 71  bbox_ = copy.deepcopy(y_arr_albu)72  pil_img = tf.keras.utils.array_to_img(X_arr)  73 74  tf.keras.utils.save_img(new_img_path, pil_img)75 76  if len(bbox_) != 0:77    bbox_[0] *= shape_after[1]78    bbox_[1] *= shape_after[0]79    bbox_[2] *= shape_after[1]80    bbox_[3] *= shape_after[0]81 82    start_point = (int(bbox_[0]), int(bbox_[1]))83    end_point = (int(bbox_[2]), int(bbox_[3]))84 85  color = (255,0,0)  86  thickness = 187 88  image = cv2.imread('./' + new_img_path)89 90  if len(bbox_) != 0:91    image = cv2.rectangle(image, start_point, end_point, color, thickness) 92    image = cv2.putText(image, 'sample ' + str(confidence), (int(bbox_[0]), int(bbox_[1] - 3)), 0, 0.3, (255,0,0), 1)93 94  final_img = tf.keras.utils.array_to_img(image)95 96  return final_img97 98 99 100def sample_detection(image):101 102    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)103    img_arr = cv2.resize(gray,(512,360)) 104    img_normalized = cv2.normalize(img_arr, None, -1, 1, cv2.NORM_MINMAX, dtype=cv2.CV_32F)105 106    img_expand = np.expand_dims(img_normalized, axis=0)107    y_pred_detection = model_detection.predict(img_expand)[0]108    confidence = round(y_pred_detection[1], 3)109    110    if confidence >= 0.3:111 112      img_resized = cv2.resize(img_normalized,(512,384))113      img_expand = np.expand_dims(img_resized, axis=0)114      y_pred_regression = model_regression.predict(img_expand)[0]115      image_with_bbox = show_image_bbox(img_resized, y_pred_regression, confidence, 'image.png')116 117      return image_with_bbox, 'There is a sample in this image'118 119    else:120 121      image = '3796506.jpg'122      return image, 'There is no sample in this image'123 124 125demo = gr.Interface(fn=sample_detection, inputs='image', outputs=['image', 'text'])126demo.launch(debug=True)