YACWC

02_remove_without_nuggets/remove_without_nuggets.py

@@ -0,0 +1,110 @@
+#!/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)
+    
+# %%
+