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))
def generate_action_model_frame(source):
CAM_SOURCE = source
class ActionModel:
def __init__(self) -> None:
self.ACTION_LIST = []
# 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)
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.pose_model = SPPE_FastPose(POSE_BACKBONE, INP_POSE[0], INP_POSE[1], device=DEVICE, path=SPPE_WEIGHT_FILE)
self.action_model = TSSTG(weight_file=TSSTG_WEIGHT_FILE) # action model
# Tracker.
max_age = 30
tracker = Tracker(max_age=max_age, n_init=3)
self.max_age = 30
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,
preprocess=preproc).start()
@ -125,6 +134,9 @@ def generate_action_model_frame(source):
elif action_name == 'Lying Down':
clr = (255, 200, 0)
# Add action to action list.
self.ACTION_LIST.append(action)
# VISUALIZE.
if track.time_since_update == 0:
if SHOW_SKELETON:
@ -149,7 +161,3 @@ def generate_action_model_frame(source):
raise HTTPException(status_code=500, detail="Frame encoding failed")
yield (b'--frame\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 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 = {
@ -24,6 +24,8 @@ jobstores = {
}
scheduler = AsyncIOScheduler(jobstores=jobstores, timezone='Asia/Bangkok')
action_model = ActionModel()
@asynccontextmanager
async def lifespan(application: FastAPI):
scheduler.start()
@ -127,7 +129,7 @@ async def stream_action_video(camera_id: int) -> StreamingResponse:
if not cap.isOpened():
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)