GrpC_Identikit/Saving/opencv dnn.py

34 lines
1.5 KiB
Python

import cv2 as cv
import numpy as np
# Chargement du modèle Mask R-CNN
net = cv.dnn.readNetFromTensorflow('frozen_inference_graph.pb', 'mask_rcnn_inception_v2_coco_2018_01_28.pbtxt')
# Charger l'image
imagepath = "photo.jpg"
image = cv.imread(imagepath)
h, w = image.shape[:2]
# Prétraiter l'image pour Mask R-CNN
blob = cv.dnn.blobFromImage(image, 1.0, (w, h), (104.0, 177.0, 123.0), swapRB=True, crop=False)
net.setInput(blob)
# Obtenir les sorties du modèle
output_layers = net.getUnconnectedOutLayersNames()
detections = net.forward(output_layers)
# Appliquer la segmentation pour la personne
mask_image = image.copy()
for detection in detections:
for obj in detection:
scores = obj[5:]
class_id = np.argmax(scores)
confidence = scores[class_id]
if class_id == 0 and confidence > 0.5: # Class 0 corresponds to "person"
# Coordonner la boîte englobante
box = obj[0:4] * np.array([w, h, w, h])
(x, y, x2, y2) = box.astype("int")
# Créer un masque de la personne
mask = np.zeros((h, w), dtype=np.uint8)
mask[y:y2, x:x2] = 255 # Définir la zone de la personne
# Appliquer le masque sur l'image originale
result = cv.bitwise_and(image, image, mask=mask)
# Montrer l'image avec la personne segmentée et l'arrière-plan supprimé
cv.imshow("Segmented Image", result)
cv.waitKey(0)
cv.destroyAllWindows()