app_.py 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125
  1. #!/usr/bin/env python
  2. from importlib import import_module
  3. import os
  4. from flask import Flask, render_template, Response
  5. import cv2
  6. import torch
  7. from utils.torch_utils import *
  8. from utils.general import check_img_size, non_max_suppression, scale_coords
  9. from models.experimental import attempt_load
  10. from models.yolo import Model
  11. from utils.datasets import LoadStreams, LoadImages
  12. import random
  13. from utils.plots import plot_one_box
  14. app = Flask(__name__)
  15. @app.route('/')
  16. def index():
  17. """Video streaming home page."""
  18. return render_template('index.html')
  19. def stream():
  20. # WEIGHT = 'runs/train/yolov7-bsw2/weights/best.pt'
  21. WEIGHT = '/root/Public/model_ssai_fine_20221209/ipark_1208_1305.pt'
  22. # model = torch.load_state_dict(WEIGHT, map_location=select_device('0'))
  23. model = attempt_load(WEIGHT, map_location='cuda:0').half()
  24. # model = TracedModel(model, select_device('0'), img_size=640).half()
  25. imgsz = 320
  26. stride = int(model.stride.max()) # model stride
  27. imgsz = check_img_size(imgsz, s=stride) # check img_size
  28. device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
  29. # Run inference
  30. model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once
  31. # Get names and colors
  32. names = model.module.names if hasattr(model, 'module') else model.names
  33. colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
  34. old_img_w = old_img_h = imgsz
  35. old_img_b = 1
  36. """Video streaming generator function."""
  37. # cap = cv2.VideoCapture('rtsp://astrodom:hdci12@192.168.170.73:554/stream1')
  38. dataset = LoadStreams('rtsp://astrodom:hdci12@192.168.170.73:554/stream1', img_size=imgsz, stride=stride)
  39. while True:
  40. # ret, frame = cap.read()
  41. # if not ret:
  42. # print('error')
  43. # break
  44. for path, img, im0s, vid_cap in dataset:
  45. img = torch.from_numpy(img).to(device)
  46. img = img.half() if True else img.float() # uint8 to fp16/32
  47. img /= 255.0 # 0 - 255 to 0.0 - 1.0
  48. if img.ndimension() == 3:
  49. img = img.unsqueeze(0)
  50. # Warmup
  51. if device.type != 'cpu' and (old_img_b != img.shape[0] or old_img_h != img.shape[2] or old_img_w != img.shape[3]):
  52. old_img_b = img.shape[0]
  53. old_img_h = img.shape[2]
  54. old_img_w = img.shape[3]
  55. for i in range(3):
  56. model(img, augment=True)[0]
  57. with torch.no_grad(): # Calculating gradients would cause a GPU memory leak
  58. pred = model(img, augment=True)[0]
  59. # print(pred)
  60. # print('*' * 30)
  61. # Apply NMS
  62. pred = non_max_suppression(pred, 0.40, 0.45)
  63. # Process detections
  64. for i, det in enumerate(pred): # detections per image
  65. if True: # batch_size >= 1
  66. p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count
  67. # else:
  68. # p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)
  69. p = Path(p) # to Path
  70. # save_path = str(save_dir / p.name) # img.jpg
  71. # txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # img.txt
  72. # gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
  73. if len(det):
  74. # Rescale boxes from img_size to im0 size
  75. det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
  76. # Print results
  77. for c in det[:, -1].unique():
  78. n = (det[:, -1] == c).sum() # detections per class
  79. s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
  80. # Write results
  81. for *xyxy, conf, cls in reversed(det):
  82. # if save_txt: # Write to file
  83. # xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
  84. # line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh) # label format
  85. # with open(txt_path + '.txt', 'a') as f:
  86. # f.write(('%g ' * len(line)).rstrip() % line + '\n')
  87. if True: # Add bbox to image
  88. label = f'{names[int(cls)]} {conf:.2f}'
  89. plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=1)
  90. # x = img.permute(0,3,1,2) # (B, W, H, C) --> (B, C, W, H)
  91. image_bytes = cv2.imencode('.jpg', im0)[1].tobytes()
  92. yield (b'--frame\r\n'
  93. b'Content-Type: image/jpeg\r\n\r\n' + image_bytes + b'\r\n')
  94. @app.route('/video')
  95. def video():
  96. """Video streaming route. Put this in the src attribute of an img tag."""
  97. return Response(stream(), content_type='multipart/x-mixed-replace; boundary=frame')
  98. if __name__ == '__main__' :
  99. app.run(host='0.0.0.0', threaded=True, debug=True)