123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124 |
- from importlib import import_module
- import os
- from flask import Flask, render_template, Response
- import cv2
- import torch
- from utils.torch_utils import *
- from utils.general import check_img_size, non_max_suppression, scale_coords
- from models.experimental import attempt_load
- from models.yolo import Model
- from utils.datasets import LoadStreams, LoadImages
- import random
- from utils.plots import plot_one_box
- app = Flask(__name__)
- @app.route('/')
- def index():
- """Video streaming home page."""
- return render_template('index.html')
- def stream():
-
- WEIGHT = '/root/Public/model_ssai_fine_20221209/ipark_1208_1305.pt'
-
- model = attempt_load(WEIGHT, map_location='cuda:0').half()
-
- imgsz = 640
- stride = int(model.stride.max())
- imgsz = check_img_size(imgsz, s=stride)
- device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
-
- model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))
-
- names = model.module.names if hasattr(model, 'module') else model.names
- colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]
- old_img_w = old_img_h = imgsz
- old_img_b = 1
- """Video streaming generator function."""
-
- dataset = LoadStreams('rtsp://astrodom:hdci12@192.168.170.73:554/stream1', img_size=imgsz, stride=stride)
- while True:
-
-
-
-
- for path, img, im0s, vid_cap in dataset:
- img = torch.from_numpy(img).to(device)
- img = img.half() if True else img.float()
- img /= 255.0
- if img.ndimension() == 3:
- img = img.unsqueeze(0)
-
- 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]):
- old_img_b = img.shape[0]
- old_img_h = img.shape[2]
- old_img_w = img.shape[3]
- for i in range(3):
- model(img, augment=True)[0]
- with torch.no_grad():
- pred = model(img, augment=True)[0]
-
-
-
- pred = non_max_suppression(pred, 0.40, 0.45)
-
- for i, det in enumerate(pred):
- if True:
- p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count
-
-
- p = Path(p)
-
-
-
- if len(det):
-
- det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
-
- for c in det[:, -1].unique():
- n = (det[:, -1] == c).sum()
- s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "
-
- for *xyxy, conf, cls in reversed(det):
-
-
-
-
-
- if True:
- label = f'{names[int(cls)]} {conf:.2f}'
- plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=1)
-
- image_bytes = cv2.imencode('.jpg', im0)[1].tobytes()
- yield (b'--frame\r\n'
- b'Content-Type: image/jpeg\r\n\r\n' + image_bytes + b'\r\n')
- @app.route('/video')
- def video():
- """Video streaming route. Put this in the src attribute of an img tag."""
- return Response(stream(), content_type='multipart/x-mixed-replace; boundary=frame')
- if __name__ == '__main__' :
- app.run(host='0.0.0.0', threaded=True, debug=True)
|