from importlib import import_module import os from flask import Flask, render_template, Response, request, send_file import cv2 import subprocess import time import sys import logging import pdb sys.path.append('./') # to run '$ python *.py' files in subdirectories logger = logging.getLogger(__name__) import torch from models.common import * from models.experimental import * from models.yolo import * from utils.autoanchor import check_anchor_order from utils.general import make_divisible, check_file, set_logging from utils.torch_utils import time_synchronized, fuse_conv_and_bn, model_info, scale_img, initialize_weights, \ select_device, copy_attr from utils.loss import SigmoidBin # import camera driver # from object_detection import VideoStreaming # if os.environ.get('CAMERA'): # Camera = import_module('camera_' + os.environ['CAMERA']).Camera # else: # from camera import Camera app = Flask(__name__) # def gen(camera): # while True: # frame = VideoStreaming.get_frame() # # cv2.imencode('.jpg', frame) # yield (b'--frame\r\n' # b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n') PATH_WEIGHT = './models/best.pt' img_mean = np.array([104., 117., 123.])[:, np.newaxis, np.newaxis].astype('float32') class Net(): def __init__(self, device='cuda'): tstamp = time.time() self.device = select_device(device) print('[yolo] loading with', self.device) self.net = Model('/root/helmet_det/yolov7-main/cfg/training/yolov7.yaml').to(self.device) state_dict = torch.load(PATH_WEIGHT, map_location=self.device)['model'].state_dict() self.net.load_state_dict(state_dict) self.net.eval() print('[yolo] finished loading (%.4f sec)' % (time.time() - tstamp)) def detect_faces(self, image, conf_th=0.8, scales=[1]): print(image, image.shape) print(len(image), len(image.shape)) print('*'*30) w, h = image.shape[1], image.shape[0] bboxes = np.empty(shape=(0, 5)) with torch.no_grad(): for s in scales: scaled_img = cv2.resize(image, dsize=(0, 0), fx=s, fy=s, interpolation=cv2.INTER_LINEAR) scaled_img = np.swapaxes(scaled_img, 1, 2) scaled_img = np.swapaxes(scaled_img, 1, 0) scaled_img = scaled_img[[2, 1, 0], :, :] scaled_img = scaled_img.astype('float32') scaled_img -= img_mean scaled_img = scaled_img[[2, 1, 0], :, :] x = torch.from_numpy(scaled_img).unsqueeze(0).to(self.device) y = self.net(x) detections = y.data scale = torch.Tensor([w, h, w, h]) for i in range(detections.size(1)): j = 0 while detections[0, i, j, 0] > conf_th: score = detections[0, i, j, 0] pt = (detections[0, i, j, 1:] * scale).cpu().numpy() bbox = (pt[0], pt[1], pt[2], pt[3], score) bboxes = np.vstack((bboxes, bbox)) j += 1 # keep = nms_(bboxes, 0.1) ## nms? # bboxes = bboxes[keep] return bboxes def get_stream_video(): # camera 정의 cam = cv2.VideoCapture('rtsp://astrodom:hdci12@192.168.170.73:554/stream1') model = Net() while True: # 카메라 값 불러오기 success, frame = cam.read() # print(frame) # print(type(frame)) if not success: break else: ret, buffer = cv2.imencode('.jpeg', frame) # frame을 byte로 변경 후 특정 식??으로 변환 후에 # yield로 하나씩 넘겨준다. decode_img = cv2.imdecode(buffer, 1) frame = model.detect_faces(decode_img).tobytes() # frame = buffer.tobytes() print(type(frame)) pdb.set_trace() yield (b'--frame\r\n' b'Content-Type: image/jpeg\r\n\r\n' + bytearray(frame) + b'\r\n') @app.route('/') def index(): return render_template('index.html') # 스트리밍 경로를 /video 경로로 설정. @app.get("/video") def video(): # StringResponse함수를 return하고, # 인자로 OpenCV에서 가져온 "바이트"이미지와 type을 명시 return Response(get_stream_video(), mimetype="multipart/x-mixed-replace; boundary=frame") # ipcam 열기 @app.route("/stream", methods=['GET']) def stream(): print("here") subprocess.run(['python3', '/root/helmet_det/yolov7-main/detect.py', '--source', 'rtsp://astrodom:hdci12@192.168.170.73:554/stream1', '--weights', 'best.pt']) return "done" # rtsp://astrodom:hdci12@192.168.170.73:554/stream1 if __name__ == '__main__': app.run(host='0.0.0.0', port=5000, debug=True)