def detect(video_id): """ This is a HACK since Tensorflow is absolutely atrocious in allocating and freeing up memory. Once a process / session is allocated a memory it cannot be forced to clear it up. As a result this code gets called via a subprocess which clears memory when it exits. :param video_id: :return: """ import django from PIL import Image sys.path.append(os.path.dirname(__file__)) os.environ.setdefault("DJANGO_SETTINGS_MODULE", "dva.settings") django.setup() from django.conf import settings from dvaapp.models import Video, Detection, Frame from dvalib import entity, detector dv = Video.objects.get(id=video_id) frames = Frame.objects.all().filter(video=dv) v = entity.WVideo(dvideo=dv, media_dir=settings.MEDIA_ROOT) wframes = { df.pk: entity.WFrame(video=v, frame_index=df.frame_index, primary_key=df.pk) for df in frames } detection_count = 0 detector_list = { 'ssd': detector.SSDetector(), } if 'YOLO_ENABLE' in os.environ: detector_list['yolo'] = detector.YOLODetector() for alogrithm in detector_list.itervalues(): logging.info("starting detection {}".format(alogrithm.name)) frame_detections = alogrithm.detect(wframes.values()) for frame_pk, detections in frame_detections.iteritems(): for d in detections: dd = Detection() dd.video = dv dd.frame_id = frame_pk dd.object_name = d['name'] dd.confidence = d['confidence'] dd.x = d['left'] dd.y = d['top'] dd.w = d['right'] - d['left'] dd.h = d['bot'] - d['top'] dd.save() img = Image.open(wframes[frame_pk].local_path()) img2 = img.crop((d['left'], d['top'], d['right'], d['bot'])) img2.save("{}/{}/detections/{}.jpg".format( settings.MEDIA_ROOT, video_id, dd.pk)) detection_count += 1 dv.refresh_from_db() dv.detections = dv.detections + detection_count dv.save()
def load_detector(self,cd): if cd.pk not in DetectorTask._detectors: if cd.detector_type == DeepModel.TFD: DetectorTask._detectors[cd.pk] = detector.TFDetector(model_path=cd.get_model_path(), class_index_to_string=cd.class_index_to_string) elif cd.detector_type == DeepModel.YOLO: DetectorTask._detectors[cd.pk] = detector.YOLODetector(cd.get_yolo_args()) elif cd.name == 'face': DetectorTask._detectors[cd.pk] = detector.FaceDetector() elif cd.name == 'textbox': DetectorTask._detectors[cd.pk] = detector.TextBoxDetector(model_path=cd.get_model_path()) else: raise ValueError,"{}".format(cd.pk)
def load_detector(cls, cd): cd.ensure() if cd.pk not in Detectors._detectors: if cd.detector_type == TrainedModel.TFD: Detectors._detectors[cd.pk] = detector.TFDetector( model_path=cd.get_model_path(), class_index_to_string=cd.arguments['class_index_to_string'] ) elif cd.detector_type == TrainedModel.YOLO: Detectors._detectors[cd.pk] = detector.YOLODetector( cd.get_yolo_args()) elif cd.name == 'face': Detectors._detectors[cd.pk] = detector.FaceDetector() elif cd.name == 'textbox': Detectors._detectors[cd.pk] = detector.TextBoxDetector( model_path=cd.get_model_path()) else: raise ValueError, "{}".format(cd.pk)
def load_detector(cls, cd): cd.ensure() if cd.pk not in Detectors._detectors: if cd.detector_type == TrainedModel.TFD: Detectors._detectors[cd.pk] = detector.TFDetector( model_path=cd.get_model_path(), class_index_to_string=cd.arguments['class_index_to_string'] ) elif cd.detector_type == TrainedModel.YOLO: # class_names = {k: v for k, v in json.loads(self.class_names)} # args = {'root_dir': model_dir, # 'detector_pk': self.pk, # 'class_names':{i: k for k, i in class_names.items()} # } Detectors._detectors[cd.pk] = detector.YOLODetector( cd.get_yolo_args()) elif cd.name == 'face': Detectors._detectors[cd.pk] = detector.FaceDetector() elif cd.name == 'textbox': Detectors._detectors[cd.pk] = detector.TextBoxDetector( model_path=cd.get_model_path()) else: raise ValueError, "{}".format(cd.pk)
def load_detector(self, detector_pk, i_class_names, model_dir): argdict = {'root_dir': model_dir, 'detector_pk': detector_pk} DetectorTask._detectors[detector_pk] = detector.YOLODetector( i_class_names=i_class_names, args=argdict)