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)
