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