Spaces:
Runtime error
Runtime error
File size: 2,168 Bytes
7a3a9cc |
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 |
import os
from pathlib import Path
import numpy as np
import cv2
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.utils.visualizer import ColorMode
PATH_PROJECT = Path(__file__).parent.parent
def get_model():
"""
This function is for the model of the project
"""
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 # set threshold for this model
cfg.MODEL.WEIGHTS = str(PATH_PROJECT/"output"/"model_final.pth") # Let training initialize from model zoo
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2
predictor = DefaultPredictor(cfg)
return predictor
def predict_image(img_pil):
"""
This function is for the prediction of the model
return the image with the prediction and the areas of the objects
"""
predictor = get_model()
img_array = np.array(img_pil)
outputs = predictor(img_array)
v = Visualizer(img_array,
scale=1,
instance_mode=ColorMode.IMAGE_BW # remove the colors of unsegmented pixels
)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
image_output = v.get_image()[:, :, ::-1].copy()
masks = outputs["instances"].pred_masks.cpu().numpy()
class_ids = outputs["instances"].pred_classes.cpu().numpy()
areas = [np.sum(mask) for mask in masks]
# Add text labels with the object IDs and areas
for i, (mask, class_id, area) in enumerate(zip(masks, class_ids, areas)):
text = f"The id is {i}"
text_size, _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, thickness=1)
pos = (np.unravel_index(np.argmax(mask), mask.shape))[::-1]
pos = (pos[0] - text_size[0]//2, pos[1] - text_size[1]//2)
cv2.putText(image_output, text, pos, cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, color=(0,0,255), thickness=1)
values = {"image":image_output, "areas":areas, "masks":masks}
return values
|