drsaikirant88's picture
initial commit with working code (local)
b202543
raw
history blame
4.53 kB
# PyTorch implementation of Darknet
# This is a custom, hard-coded version of darknet with
# YOLOv3 implementation for openimages database. This
# was written to test viability of implementing YOLO
# for face detection followed by emotion / sentiment
# analysis.
#
# Configuration, weights and data are hardcoded.
# Additional options include, ability to create
# subset of data with faces exracted for labelling.
#
# Author : Saikiran Tharimena
# Co-Authors: Kjetil Marinius Sjulsen, Juan Carlos Calvet Lopez
# Project : Emotion / Sentiment Detection from news images
# Date : 12 September 2022
# Version : v0.1
#
# (C) Schibsted ASA
# Libraries
import os
import cv2
import torch
import numpy as np
from utils import *
from darknet import Darknet
from torch.autograd import Variable
from torch.cuda import is_available as check_cuda
# Parameters
batch_size = 1
confidence = 0.25
nms_thresh = 0.30
run_cuda = False
# CFG Files
cwd = os.path.dirname(__file__)
cfg = cwd + '/cfg/yolov3-openimages.cfg'
data = cwd + '/cfg/openimages.data'
clsnames= cwd + '/cfg/openimages.names'
weights = cwd + '/cfg/yolov3-openimages.weights'
# Load classes
num_classes = 601
classes = load_classes(clsnames)
# Set up the neural network
print('Load Network')
model = Darknet(cfg)
print('Load Weights')
model.load_weights(weights)
print('Successfully loaded Network')
# Check CUDA
if run_cuda:
CUDA = check_cuda()
else:
CUDA = False
# Input dimension
inp_dim = int(model.net_info["height"])
# put the model on GPU
if CUDA:
model.cuda()
# Set the model in evaluation mode
model.eval()
# face detector
def detect_face(image):
# Just lazy to update this
imlist = [image]
loaded_ims = [cv2.imread(x) for x in imlist]
im_batches = list(map(prep_image, loaded_ims, [inp_dim for x in range(len(imlist))]))
im_dim_list = [(x.shape[1], x.shape[0]) for x in loaded_ims]
im_dim_list = torch.FloatTensor(im_dim_list).repeat(1,2)
leftover = 0
if (len(im_dim_list) % batch_size):
leftover = 1
if batch_size != 1:
num_batches = len(imlist) // batch_size + leftover
im_batches = [torch.cat((im_batches[i*batch_size : min((i + 1)*batch_size,
len(im_batches))])) for i in range(num_batches)]
write = 0
if CUDA:
im_dim_list = im_dim_list.cuda()
for i, batch in enumerate(im_batches):
# load the image
if CUDA:
batch = batch.cuda()
with torch.no_grad():
prediction = model(Variable(batch), CUDA)
prediction = write_results(prediction, confidence, num_classes, nms_conf = nms_thresh)
if type(prediction) == int:
for im_num, image in enumerate(imlist[i*batch_size: min((i + 1)*batch_size, len(imlist))]):
im_id = i*batch_size + im_num
continue
prediction[:,0] += i*batch_size # transform the atribute from index in batch to index in imlist
if not write: # If we have't initialised output
output = prediction
write = 1
else:
output = torch.cat((output, prediction))
for im_num, image in enumerate(imlist[i*batch_size: min((i + 1)*batch_size, len(imlist))]):
im_id = i * batch_size + im_num
objs = [classes[int(x[-1])] for x in output if int(x[0]) == im_id]
if CUDA:
torch.cuda.synchronize()
try:
output
except NameError:
return None
im_dim_list = torch.index_select(im_dim_list, 0, output[:,0].long())
scaling_factor = torch.min(608/im_dim_list,1)[0].view(-1,1)
output[:, [1,3]] -= (inp_dim - scaling_factor*im_dim_list[:,0].view(-1,1))/2
output[:, [2,4]] -= (inp_dim - scaling_factor*im_dim_list[:,1].view(-1,1))/2
output[:, 1:5] /= scaling_factor
for i in range(output.shape[0]):
output[i, [1,3]] = torch.clamp(output[i, [1,3]], 0.0, im_dim_list[i,0])
output[i, [2,4]] = torch.clamp(output[i, [2,4]], 0.0, im_dim_list[i,1])
def get_detections(x, results):
c1 = [int(y) for y in x[1:3]]
c2 = [int(y) for y in x[3:5]]
det_class = int(x[-1])
label = "{0}".format(classes[det_class])
return (label, tuple(c1 + c2))
detections = list(map(lambda x: get_detections(x, loaded_ims), output))
if CUDA:
torch.cuda.empty_cache()
return loaded_ims[0], detections