GrpC_Identikit/ImageProcessing/Tensorflow.py

30 lines
1.2 KiB
Python

import tensorflow as tf
import numpy as np
import cv2
import os
# Charger le modèle Mask R-CNN
model = tf.saved_model.load('frozen_inference_graph.pb')
# Charger l'image
image = cv2.imread('photo.jpg')
# Prétraiter l'image
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
input_tensor = tf.convert_to_tensor(image_rgb)
input_tensor = input_tensor[tf.newaxis,...] # Ajouter une dimension batch
# Effectuer la détection
detections = model(input_tensor)
# Extraire les masques des objets détectés
masks = detections['detection_masks'][0].numpy()
boxes = detections['detection_boxes'][0].numpy()
class_ids = detections['detection_classes'][0].numpy()
# Sélectionner uniquement les personnes (classe 1 dans COCO)
for i in range(len(masks)):
if class_ids[i] == 1: # Personne
mask = masks[i]
mask = mask > 0.5 # Seuil pour la segmentation binaire
# Appliquer le masque sur l'image d'origine
mask = np.uint8(mask * 255) # Convertir en format compatible OpenCV
result = cv2.bitwise_and(image, image, mask=mask)
# Afficher l'image segmentée
cv2.imshow("Segmented Image", result)
cv2.waitKey(0)
cv2.destroyAllWindows()