Tryag File Manager

//paip/script/distribution_abnormal/person_detecting.py

# Copyright (c) 2020~2023 PAIPTREE SmartFarm, Inc. All Rights Reserved.
# Author: 박정훈
# modified : 2024-04-04
# GPU/CPU 실행 로직부분 버그수정

import torch
import sys, glob

# Some basic setup:
# Setup detectron2 logger
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()

# import some common libraries
import os, cv2
import tensorflow as tf

# import some common detectron2 utilities
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog, DatasetCatalog

if __name__ == '__main__':
    if len(sys.argv) <= 1 :
        print("fail - not input img_path")
        sys.exit()

print("Start to Detect person!! ")

img_path = sys.argv[1]
    print("Input img_path : ", img_path)

try :
        # find model file
        windows_model_path = './weight/'
        linux_model_path = os.path.dirname(os.path.abspath(__file__)) + '/weight/'
        model_path = windows_model_path if os.name == 'nt' else linux_model_path
        #print("model_path : ", model_path)

model_file_list = glob.glob(model_path + '*')
        model_file_path = max(model_file_list, key=os.path.getmtime)
        #print("model_file_path : ", model_file_path)
        
        #수정/주석/삭제 금지 : gateway에서 이 구문을 통해 파일명 획득함
        print("model_file_name:", os.path.basename(model_file_path))

cfg = get_cfg()
        
        if len(tf.config.experimental.list_physical_devices('GPU')) <= 0 :
            print('Running CPU!!')
            cfg.merge_from_list(['MODEL.DEVICE', 'cpu'])  # CPU 변경
        
        # add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library
        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.5  # set threshold for this model
        cfg.MODEL.WEIGHTS = model_path + "model_final_f10217.pkl"

predictor = DefaultPredictor(cfg)
        
        im = cv2.imread(img_path)
        outputs = predictor(im)     
        output_cpu = outputs['instances'].to("cpu")
        # print(outputs["instances"].pred_classes)
        # print(outputs["instances"].pred_boxes)
        output_list = output_cpu.pred_classes.tolist()

if 0 in output_list:
            print("Yes Person!\n")
            
            toggle = [i for i in range(len(output_list)) if 0 == output_list[i]]
            print('score:%s' % output_cpu[toggle].scores.tolist())
            
            v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.0)
            out = v.overlay_instances(boxes=output_cpu[toggle].pred_boxes, assigned_colors='red') 
            # out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
            # cv2_imshow(out.get_image()[:, :, ::-1])
            
            result_image = img_path.split('.')[0] + '_personDetection.jpg'
            cv2.imwrite(result_image, out.get_image()[:, :, ::-1])
            print('1:' + os.path.basename(result_image))
        else:
            print("No Person!\n")
            print('0')

except :
        print("fail")