111 lines
3.3 KiB
Python
Executable File
111 lines
3.3 KiB
Python
Executable File
#!/usr/bin/python
|
|
import glob
|
|
import argparse
|
|
import os
|
|
import json
|
|
from CommonCode.util import exit_if_not_ipython, get_cset_for_file_matching, is_ipython
|
|
from CommonCode import kwq
|
|
from CommonCode.settings import get_logger, LogColorize
|
|
from kafka import TopicPartition
|
|
from kafka.structs import OffsetAndMetadata
|
|
from pprint import pprint
|
|
|
|
pfm = LogColorize.remove_without_nuggets
|
|
logger = get_logger(__name__,'/var/log/ml_vision_logs/02_remove_without_nuggets', stdout=True, systemd=False)
|
|
|
|
input_topic = kwq.TOPICS.videos_scored_detection
|
|
|
|
client_id='obj_detector_remove_without_nuggets'
|
|
group_id = client_id
|
|
|
|
consumer = kwq.create_consumer(input_topic, group_id = group_id, client_id = client_id)
|
|
c_part = TopicPartition(input_topic, 0)
|
|
consumer.assign([c_part])
|
|
|
|
producer = kwq.producer
|
|
|
|
|
|
class RESULT_TYPE():
|
|
NOTHING = 0
|
|
HAS_OBJS = 1
|
|
KEPT = 2
|
|
REMOVED = 3
|
|
NO_JSON = 4
|
|
|
|
|
|
def exec_file_remove_logic(file_path):
|
|
result = RESULT_TYPE.NOTHING
|
|
cset = get_cset_for_file_matching(file_path)
|
|
logger.info(f"EXECUTING_LOGIC :{file_path}")
|
|
json_check = '.json.orin'
|
|
if '.has_objs' in cset:
|
|
logger.info(f"HAS_OBJS :{file_path}")
|
|
result = RESULT_TYPE.HAS_OBJS
|
|
return result
|
|
|
|
|
|
if json_check in cset and os.path.exists(cset[ json_check]):
|
|
logger.info(f"HAS_JSON :{file_path}")
|
|
if os.path.getsize(cset[json_check]) == 0:
|
|
det_data = {'scores':[]}
|
|
else:
|
|
with open(cset[json_check],'r') as jj:
|
|
det_data = json.load(jj)
|
|
|
|
thresh = 0.2
|
|
frames_with_dets_thresh = 0.1
|
|
|
|
all_scores = list()
|
|
for x in det_data['scores']:
|
|
sc_in = [y['score'] for y in x['detections']]
|
|
sc_in.append(0)
|
|
all_scores.append(max(sc_in))
|
|
|
|
frames_above_thresh = sum([x>thresh for x in all_scores ])
|
|
num_frames = len(det_data['scores'])
|
|
|
|
ratio_with_frames = frames_above_thresh / (num_frames + 1)
|
|
if ratio_with_frames > frames_with_dets_thresh:
|
|
cpath = os.path.splitext(cset['.mp4'])[0] + '.has_objs.orin'
|
|
with open(cpath,'w') as cc:
|
|
pass
|
|
|
|
|
|
logger.info(f"HAS_OBJECTS_DETECTED_WILL_NOT_REMOVE :{pfm(file_path)}")
|
|
result = RESULT_TYPE.KEPT
|
|
|
|
else:
|
|
logger.info(f"OBJECTS_NOT_DETECTED_WILL_REMOVE :{pfm(file_path)}")
|
|
for ff, cf in cset.items():
|
|
if cf.endswith(json_check) or cf.endswith('.jpg'):
|
|
pass
|
|
else:
|
|
logger.info(f"REMOVING :{cf}")
|
|
os.remove(cf)
|
|
|
|
result = RESULT_TYPE.REMOVED
|
|
|
|
else:
|
|
result = RESULT_TYPE.NO_JSON
|
|
logger.info(f"NO_JSON_DONT_REMOVE :{pfm(file_path)}")
|
|
|
|
|
|
return result
|
|
|
|
|
|
for msg in consumer:
|
|
key = msg.key
|
|
value = msg.value
|
|
|
|
purge_reason = exec_file_remove_logic(value['filepath'])
|
|
d_send = {'value':msg.value, 'key':msg.key}
|
|
if purge_reason in [RESULT_TYPE.KEPT, RESULT_TYPE.HAS_OBJS]:
|
|
producer.send(kwq.TOPICS.videos_with_nuggets, **d_send)
|
|
elif purge_reason == RESULT_TYPE.REMOVED:
|
|
producer.send(kwq.TOPICS.videos_without_nuggets, **d_send)
|
|
elif purge_reason == RESULT_TYPE.NO_JSON:
|
|
producer.send(kwq.TOPICS.videos_no_json, **d_send)
|
|
|
|
# %%
|
|
|