Spaces:
Running
Running
File size: 2,838 Bytes
9d52aa6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
from transformers import pipeline
from pathlib import Path
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import colormaps
from matplotlib.patches import Patch
import cv2
from matplotlib import patches
detc_modelo = 'valentinafeve/yolos-fashionpedia'
detc = pipeline('object-detection', detc_modelo)
sgm_modelo = 'mattmdjaga/segformer_b0_clothes'
sgm = pipeline('image-segmentation', sgm_modelo)
def processar_imagem(image_path):
# Carrega a imagem
imagem = Image.open(Path(image_path))
# Segmentação
segmentacao = sgm(imagem)
# Detecção de objetos
deteccao = detc(imagem)
return imagem, segmentacao, deteccao
# segmentacao = sgm(imagem) feito na função processar_imagem
def adicionar_mascara(dados_imagem, dados_mascara, cor):
imagem_mascara = dados_imagem.copy()
imagem_mascara = np.where(
dados_mascara[:, :, np.newaxis] == 255,
cor,
imagem_mascara,
).astype(dados_imagem.dtype)
return cv2.addWeighted(dados_imagem, 0.5, imagem_mascara, 0.5, 0)
def plotar_segmentos(imagem, segmentacao, nome_colormap):
cmap = colormaps.get_cmap(nome_colormap)
cores = [
(np.array(cmap(x)[:3]) * 255).astype(int)
for x in np.linspace(0, 1, len(segmentacao))
]
imagem_final = np.array(imagem).copy()
legendas = []
for segmento, cor in zip(segmentacao, cores):
dados_mascara = np.array(segmento['mask'])
label_mascara = segmento['label']
imagem_final = adicionar_mascara(imagem_final, dados_mascara, cor)
legendas.append(Patch(facecolor=cor/255, edgecolor='black', label=label_mascara))
fig, ax = plt.subplots(figsize=(16, 16))
plt.legend(handles=legendas)
ax.imshow(imagem_final)
#plotar_segmentos(imagem=imagem, segmentacao=segmentacao, nome_colormap='hsv')
# deteccao = detc(imagem) feito na função processar_imagem
# Função para desenhar os retângulos
def plotar_deteccao(imagem, deteccao):
fig, ax = plt.subplots(figsize=(16, 16))
ax.imshow(imagem)
for det in deteccao:
if det['score'] >= 0.80:
box = det['box']
origem = (box['xmin'], box['ymin'])
largura = box['xmax'] - box['xmin']
altura = box['ymax'] - box['ymin']
rect = patches.Rectangle(origem, largura, altura, linewidth=3, edgecolor='red', facecolor='none')
ax.add_patch(rect)
texto = f"{det['label']} {100 * det['score']:.2f}%"
ax.text(box['xmin'] - 50, box['ymin'] + altura / 2, texto,
bbox={'facecolor': 'red', 'alpha': 0.8}, ha='right', va='center')
plt.axis('off')
plt.show()
# plotar_deteccao(imagem=imagem, deteccao=deteccao) |