def yolo_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,Region,Frame from dvaapp.operations.video_processing import WVideo,WFrame from dvalib import detector dv = Video.objects.get(id=video_id) frames = Frame.objects.all().filter(video=dv) v = WVideo(dvideo=dv, media_dir=settings.MEDIA_ROOT) wframes = {df.pk: WFrame(video=v, frame_index=df.frame_index, primary_key=df.pk) for df in frames} detection_count = 0 algorithm = detector.YOLODetector() logging.info("starting detection {}".format(algorithm.name)) frame_detections = algorithm.detect(wframes.values()) for frame_pk,detections in frame_detections.iteritems(): for d in detections: dd = Region() dd.region_type = Region.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 detect_custom_objects(detector_pk,video_pk): """ Detection using customized trained YOLO detectors :param detector_pk: :param video_pk: :return: """ setup_django() from dvaapp.models import Region, Frame, CustomDetector from django.conf import settings from dvalib.yolo import trainer from PIL import Image args = {'detector_pk':int(detector_pk)} video_pk = int(video_pk) detector = CustomDetector.objects.get(pk=args['detector_pk']) args['root_dir'] = "{}/models/{}/".format(settings.MEDIA_ROOT, detector.pk) class_names = {k:v for k,v in json.loads(detector.class_names)} i_class_names = {i: k for k, i in class_names.items()} frames = {} for f in Frame.objects.all().filter(video_id=video_pk): frames[f.pk] = f images = [] path_to_f = {} for k,f in frames.iteritems(): path = "{}/{}/frames/{}.jpg".format(settings.MEDIA_ROOT,f.video_id,f.frame_index) path_to_f[path] = f images.append(path) train_task = trainer.YOLOTrainer(boxes=[], images=images, class_names=i_class_names, args=args,test_mode=True) results = train_task.predict() for path, box_class, score, top, left, bottom, right in results: r = Region() r.region_type = r.DETECTION r.confidence = int(100.0 * score) r.object_name = "YOLO_{}_{}".format(detector.pk, box_class) r.y = top r.x = left r.w = right - left r.h = bottom - top r.frame_id = path_to_f[path].pk r.video_id = path_to_f[path].video_id r.save() right = r.w + r.x bottom = r.h + r.y img = Image.open(path) img2 = img.crop((r.x,r.y,right, bottom)) img2.save("{}/{}/detections/{}.jpg".format(settings.MEDIA_ROOT, video_pk, r.pk))
def detect_custom_objects(detector_pk,video_pk): """ Detection using customized trained YOLO detectors :param detector_pk: :param video_pk: :return: """ setup_django() from dvaapp.models import Region, Frame, CustomDetector from django.conf import settings from dvalib.yolo import trainer from PIL import Image args = {'detector_pk':int(detector_pk)} video_pk = int(video_pk) detector = CustomDetector.objects.get(pk=args['detector_pk']) args['root_dir'] = "{}/detectors/{}/".format(settings.MEDIA_ROOT, detector.pk) class_names = {k:v for k,v in json.loads(detector.class_names)} i_class_names = {i: k for k, i in class_names.items()} frames = {} for f in Frame.objects.all().filter(video_id=video_pk): frames[f.pk] = f images = [] path_to_f = {} for k,f in frames.iteritems(): path = "{}/{}/frames/{}.jpg".format(settings.MEDIA_ROOT,f.video_id,f.frame_index) path_to_f[path] = f images.append(path) train_task = trainer.YOLOTrainer(boxes=[], images=images, class_names=i_class_names, args=args,test_mode=True) results = train_task.predict() for path, box_class, score, top, left, bottom, right in results: r = Region() r.region_type = r.DETECTION r.confidence = int(100.0 * score) r.object_name = "YOLO_{}_{}".format(detector.pk, box_class) r.y = top r.x = left r.w = right - left r.h = bottom - top r.frame_id = path_to_f[path].pk r.video_id = path_to_f[path].video_id r.save() right = r.w + r.x bottom = r.h + r.y img = Image.open(path) img2 = img.crop((r.x,r.y,right, bottom)) img2.save("{}/{}/regions/{}.jpg".format(settings.MEDIA_ROOT, video_pk, r.pk))
def train_yolo(start_pk): """ Train a yolo model specified in a TaskEvent. This is necessary to ensure that the Tensorflow process exits and releases the allocated GPU memory. :param start_pk: TEvent PK with information about lauching the training task :return: """ setup_django() from django.conf import settings from dvaapp.models import Region, Frame, CustomDetector, TEvent from dvaapp.shared import create_detector_folders, create_detector_dataset from dvalib.yolo import trainer start = TEvent.objects.get(pk=start_pk) args = json.loads(start.arguments_json) labels = set(args['labels']) if 'labels' in args else set() object_names = set(args['object_names']) if 'object_names' in args else set() detector = CustomDetector.objects.get(pk=args['detector_pk']) create_detector_folders(detector) args['root_dir'] = "{}/detectors/{}/".format(settings.MEDIA_ROOT,detector.pk) class_distribution, class_names, rboxes, rboxes_set, frames, i_class_names = create_detector_dataset(object_names,labels) images, boxes = [], [] path_to_f = {} for k,f in frames.iteritems(): path = "{}/{}/frames/{}.jpg".format(settings.MEDIA_ROOT,f.video_id,f.frame_index) path_to_f[path] = f images.append(path) boxes.append(rboxes[k]) # print k,rboxes[k] with open("{}/input.json".format(args['root_dir']),'w') as input_data: json.dump({'boxes':boxes, 'images':images, 'args':args, 'class_names':class_names.items(), 'class_distribution':class_distribution.items()}, input_data) detector.boxes_count = sum([len(k) for k in boxes]) detector.frames_count = len(images) detector.classes_count = len(class_names) detector.save() train_task = trainer.YOLOTrainer(boxes=boxes,images=images,class_names=i_class_names,args=args) train_task.train() detector.phase_1_log = file("{}/phase_1.log".format(args['root_dir'])).read() detector.phase_2_log = file("{}/phase_2.log".format(args['root_dir'])).read() detector.class_distribution = json.dumps(class_distribution.items()) detector.class_names = json.dumps(class_names.items()) detector.trained = True detector.save() results = train_task.predict() bulk_regions = [] for path, box_class, score, top, left, bottom, right in results: r = Region() r.region_type = r.ANNOTATION r.confidence = int(100.0 * score) r.object_name = "YOLO_{}_{}".format(detector.pk,box_class) r.y = top r.x = left r.w = right - left r.h = bottom - top r.frame_id = path_to_f[path].pk r.video_id = path_to_f[path].video_id bulk_regions.append(r) Region.objects.bulk_create(bulk_regions,batch_size=1000) folder_name = "{}/detectors/{}".format(settings.MEDIA_ROOT,detector.pk) file_name = '{}/exports/{}.dva_detector.zip'.format(settings.MEDIA_ROOT,detector.pk) zipper = subprocess.Popen(['zip', file_name, '-r', '.'],cwd=folder_name) zipper.wait() return 0
train_task = trainer.YOLOTrainer(boxes=boxes, images=images, args=args) train_task.train() detector.phase_1_log = file("{}/phase_1.log".format( args['root_dir'])).read() detector.phase_2_log = file("{}/phase_2.log".format( args['root_dir'])).read() detector.class_distribution = json.dumps(class_distribution.items()) detector.class_names = json.dumps(class_names.items()) detector.trained = True detector.save() results = train_task.predict() bulk_regions = [] for path, box_class, score, top, left, bottom, right in results: r = Region() r.region_type = r.ANNOTATION r.confidence = int(100.0 * score) r.object_name = "YOLO_{}_{}".format(detector.pk, box_class) r.y = top r.x = left r.w = right - left r.h = bottom - top r.frame_id = path_to_f[path].pk r.video_id = path_to_f[path].video_id bulk_regions.append(r) Region.objects.bulk_create(bulk_regions, batch_size=1000) folder_name = "{}/detectors/{}".format(settings.MEDIA_ROOT, detector.pk) file_name = '{}/exports/{}.dva_detector.zip'.format( settings.MEDIA_ROOT, detector.pk) zipper = subprocess.Popen(['zip', file_name, '-r', '.'], cwd=folder_name) zipper.wait()
def train_yolo(start_pk): """ Train a yolo model specified in a TaskEvent. This is necessary to ensure that the Tensorflow process exits and releases the allocated GPU memory. :param start_pk: TEvent PK with information about lauching the training task :return: """ setup_django() from django.conf import settings from dvaapp.models import Region, Frame, CustomDetector, TEvent from dvaapp.shared import create_detector_folders, create_detector_dataset from dvalib.yolo import trainer start = TEvent.objects.get(pk=start_pk) args = json.loads(start.arguments_json) labels = set(args['labels']) if 'labels' in args else set() object_names = set(args['object_names']) if 'object_names' in args else set() detector = CustomDetector.objects.get(pk=args['detector_pk']) create_detector_folders(detector) args['root_dir'] = "{}/detectors/{}/".format(settings.MEDIA_ROOT,detector.pk) class_distribution, class_names, rboxes, rboxes_set, frames, i_class_names = create_detector_dataset(object_names,labels) images, boxes = [], [] path_to_f = {} for k,f in frames.iteritems(): path = "{}/{}/frames/{}.jpg".format(settings.MEDIA_ROOT,f.video_id,f.frame_index) path_to_f[path] = f images.append(path) boxes.append(rboxes[k]) # print k,rboxes[k] with open("{}/input.json".format(args['root_dir']),'w') as input_data: json.dump({'boxes':boxes, 'images':images, 'args':args, 'class_names':class_names.items(), 'class_distribution':class_distribution.items()}, input_data) detector.boxes_count = sum([len(k) for k in boxes]) detector.frames_count = len(images) detector.classes_count = len(class_names) detector.save() train_task = trainer.YOLOTrainer(boxes=boxes,images=images,class_names=i_class_names,args=args) train_task.train() detector.phase_1_log = file("{}/phase_1.log".format(args['root_dir'])).read() detector.phase_2_log = file("{}/phase_2.log".format(args['root_dir'])).read() detector.class_distribution = json.dumps(class_distribution.items()) detector.class_names = json.dumps(class_names.items()) detector.trained = True detector.save() results = train_task.predict() bulk_regions = [] for path, box_class, score, top, left, bottom, right in results: r = Region() r.region_type = r.ANNOTATION r.confidence = int(100.0 * score) r.object_name = "YOLO_{}_{}".format(detector.pk,box_class) r.y = top r.x = left r.w = right - left r.h = bottom - top r.frame_id = path_to_f[path].pk r.video_id = path_to_f[path].video_id bulk_regions.append(r) Region.objects.bulk_create(bulk_regions,batch_size=1000) folder_name = "{}/detectors/{}".format(settings.MEDIA_ROOT,detector.pk) file_name = '{}/exports/{}.dva_detector.zip'.format(settings.MEDIA_ROOT,detector.pk) zipper = subprocess.Popen(['zip', file_name, '-r', '.'],cwd=folder_name) zipper.wait() return 0