Use ActionModel class to wrap around function that generate frame

This commit is contained in:
sosokker 2024-05-09 23:07:16 +07:00
parent 754eaf27f2
commit 67bae60e3a
2 changed files with 112 additions and 102 deletions

View File

@ -49,19 +49,28 @@ def kpt2bbox(kpt, ex=20):
kpt[:, 0].max() + ex, kpt[:, 1].max() + ex)) kpt[:, 0].max() + ex, kpt[:, 1].max() + ex))
def generate_action_model_frame(source): class ActionModel:
CAM_SOURCE = source def __init__(self) -> None:
self.ACTION_LIST = []
# Model initialization # Model initialization
detect_model = TinyYOLOv3_onecls(INP_DETS, device=DEVICE, config_file=CONFIG_FILE, self.detect_model = TinyYOLOv3_onecls(INP_DETS, device=DEVICE, config_file=CONFIG_FILE,
weight_file=YOLO_WEIGHT_FILE) weight_file=YOLO_WEIGHT_FILE)
pose_model = SPPE_FastPose(POSE_BACKBONE, INP_POSE[0], INP_POSE[1], device=DEVICE, path=SPPE_WEIGHT_FILE) self.pose_model = SPPE_FastPose(POSE_BACKBONE, INP_POSE[0], INP_POSE[1], device=DEVICE, path=SPPE_WEIGHT_FILE)
action_model = TSSTG(weight_file=TSSTG_WEIGHT_FILE) # action model self.action_model = TSSTG(weight_file=TSSTG_WEIGHT_FILE) # action model
# Tracker. # Tracker.
max_age = 30 self.max_age = 30
tracker = Tracker(max_age=max_age, n_init=3) self.tracker = Tracker(max_age=self.max_age, n_init=3)
def generate_action_model_frame(self, source):
CAM_SOURCE = source
detect_model = self.detect_model
pose_model = self.pose_model
action_model = self.action_model
max_age = self.max_age
tracker = self.tracker
cam = CamLoader(int(CAM_SOURCE) if CAM_SOURCE.isdigit() else CAM_SOURCE, cam = CamLoader(int(CAM_SOURCE) if CAM_SOURCE.isdigit() else CAM_SOURCE,
preprocess=preproc).start() preprocess=preproc).start()
@ -125,6 +134,9 @@ def generate_action_model_frame(source):
elif action_name == 'Lying Down': elif action_name == 'Lying Down':
clr = (255, 200, 0) clr = (255, 200, 0)
# Add action to action list.
self.ACTION_LIST.append(action)
# VISUALIZE. # VISUALIZE.
if track.time_since_update == 0: if track.time_since_update == 0:
if SHOW_SKELETON: if SHOW_SKELETON:
@ -149,7 +161,3 @@ def generate_action_model_frame(source):
raise HTTPException(status_code=500, detail="Frame encoding failed") raise HTTPException(status_code=500, detail="Frame encoding failed")
yield (b'--frame\r\n' yield (b'--frame\r\n'
b'Content-Type: image/jpeg\r\n\r\n' + buffer.tobytes() + b'\r\n') b'Content-Type: image/jpeg\r\n\r\n' + buffer.tobytes() + b'\r\n')
def output_action_detection():
pass

View File

@ -16,7 +16,7 @@ from config import TEMP_VIDEO_FILE, VIDEO_BUCKET
from scheme import Camera from scheme import Camera
from utils import save_to_config, read_cameras_from_config from utils import save_to_config, read_cameras_from_config
from analytic.action.action_model import generate_action_model_frame from analytic.action.action_model import ActionModel
jobstores = { jobstores = {
@ -24,6 +24,8 @@ jobstores = {
} }
scheduler = AsyncIOScheduler(jobstores=jobstores, timezone='Asia/Bangkok') scheduler = AsyncIOScheduler(jobstores=jobstores, timezone='Asia/Bangkok')
action_model = ActionModel()
@asynccontextmanager @asynccontextmanager
async def lifespan(application: FastAPI): async def lifespan(application: FastAPI):
scheduler.start() scheduler.start()
@ -127,7 +129,7 @@ async def stream_action_video(camera_id: int) -> StreamingResponse:
if not cap.isOpened(): if not cap.isOpened():
raise HTTPException(status_code=404, detail="Camera is closed or not available") raise HTTPException(status_code=404, detail="Camera is closed or not available")
return StreamingResponse(generate_action_model_frame(camera.link), media_type="multipart/x-mixed-replace; boundary=frame") return StreamingResponse(action_model.generate_action_model_frame(camera.link), media_type="multipart/x-mixed-replace; boundary=frame")
@router.delete("/remove/{camera_id}", response_model=dict) @router.delete("/remove/{camera_id}", response_model=dict)