app.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147
  1. from importlib import import_module
  2. import os
  3. from flask import Flask, render_template, Response, request, send_file
  4. import cv2
  5. import subprocess
  6. import time
  7. import sys
  8. import logging
  9. import pdb
  10. sys.path.append('./') # to run '$ python *.py' files in subdirectories
  11. logger = logging.getLogger(__name__)
  12. import torch
  13. from models.common import *
  14. from models.experimental import *
  15. from models.yolo import *
  16. from utils.autoanchor import check_anchor_order
  17. from utils.general import make_divisible, check_file, set_logging
  18. from utils.torch_utils import time_synchronized, fuse_conv_and_bn, model_info, scale_img, initialize_weights, \
  19. select_device, copy_attr
  20. from utils.loss import SigmoidBin
  21. # import camera driver
  22. # from object_detection import VideoStreaming
  23. # if os.environ.get('CAMERA'):
  24. # Camera = import_module('camera_' + os.environ['CAMERA']).Camera
  25. # else:
  26. # from camera import Camera
  27. app = Flask(__name__)
  28. # def gen(camera):
  29. # while True:
  30. # frame = VideoStreaming.get_frame()
  31. # # cv2.imencode('.jpg', frame)
  32. # yield (b'--frame\r\n'
  33. # b'Content-Type: image/jpeg\r\n\r\n' + frame + b'\r\n')
  34. PATH_WEIGHT = './models/best.pt'
  35. img_mean = np.array([104., 117., 123.])[:, np.newaxis, np.newaxis].astype('float32')
  36. class Net():
  37. def __init__(self, device='cuda'):
  38. tstamp = time.time()
  39. self.device = select_device(device)
  40. print('[yolo] loading with', self.device)
  41. self.net = Model('/root/helmet_det/yolov7-main/cfg/training/yolov7.yaml').to(self.device)
  42. state_dict = torch.load(PATH_WEIGHT, map_location=self.device)['model'].state_dict()
  43. self.net.load_state_dict(state_dict)
  44. self.net.eval()
  45. print('[yolo] finished loading (%.4f sec)' % (time.time() - tstamp))
  46. def detect_faces(self, image, conf_th=0.8, scales=[1]):
  47. print(image, image.shape)
  48. print(len(image), len(image.shape))
  49. print('*'*30)
  50. w, h = image.shape[1], image.shape[0]
  51. bboxes = np.empty(shape=(0, 5))
  52. with torch.no_grad():
  53. for s in scales:
  54. scaled_img = cv2.resize(image, dsize=(0, 0), fx=s, fy=s, interpolation=cv2.INTER_LINEAR)
  55. scaled_img = np.swapaxes(scaled_img, 1, 2)
  56. scaled_img = np.swapaxes(scaled_img, 1, 0)
  57. scaled_img = scaled_img[[2, 1, 0], :, :]
  58. scaled_img = scaled_img.astype('float32')
  59. scaled_img -= img_mean
  60. scaled_img = scaled_img[[2, 1, 0], :, :]
  61. x = torch.from_numpy(scaled_img).unsqueeze(0).to(self.device)
  62. y = self.net(x)
  63. detections = y.data
  64. scale = torch.Tensor([w, h, w, h])
  65. for i in range(detections.size(1)):
  66. j = 0
  67. while detections[0, i, j, 0] > conf_th:
  68. score = detections[0, i, j, 0]
  69. pt = (detections[0, i, j, 1:] * scale).cpu().numpy()
  70. bbox = (pt[0], pt[1], pt[2], pt[3], score)
  71. bboxes = np.vstack((bboxes, bbox))
  72. j += 1
  73. # keep = nms_(bboxes, 0.1) ## nms?
  74. # bboxes = bboxes[keep]
  75. return bboxes
  76. def get_stream_video():
  77. # camera 정의
  78. cam = cv2.VideoCapture('rtsp://astrodom:hdci12@192.168.170.73:554/stream1')
  79. model = Net()
  80. while True:
  81. # 카메라 값 불러오기
  82. success, frame = cam.read()
  83. # print(frame)
  84. # print(type(frame))
  85. if not success:
  86. break
  87. else:
  88. ret, buffer = cv2.imencode('.jpeg', frame)
  89. # frame을 byte로 변경 후 특정 식??으로 변환 후에
  90. # yield로 하나씩 넘겨준다.
  91. decode_img = cv2.imdecode(buffer, 1)
  92. frame = model.detect_faces(decode_img).tobytes()
  93. # frame = buffer.tobytes()
  94. print(type(frame))
  95. pdb.set_trace()
  96. yield (b'--frame\r\n' b'Content-Type: image/jpeg\r\n\r\n' + bytearray(frame) + b'\r\n')
  97. @app.route('/')
  98. def index():
  99. return render_template('index.html')
  100. # 스트리밍 경로를 /video 경로로 설정.
  101. @app.get("/video")
  102. def video():
  103. # StringResponse함수를 return하고,
  104. # 인자로 OpenCV에서 가져온 "바이트"이미지와 type을 명시
  105. return Response(get_stream_video(), mimetype="multipart/x-mixed-replace; boundary=frame")
  106. # ipcam 열기
  107. @app.route("/stream", methods=['GET'])
  108. def stream():
  109. print("here")
  110. subprocess.run(['python3', '/root/helmet_det/yolov7-main/detect.py', '--source', 'rtsp://astrodom:hdci12@192.168.170.73:554/stream1', '--weights', 'best.pt'])
  111. return "done"
  112. # rtsp://astrodom:hdci12@192.168.170.73:554/stream1
  113. if __name__ == '__main__':
  114. app.run(host='0.0.0.0', port=5000, debug=True)