Skip to content

Instantly share code, notes, and snippets.

@dfaker
Created June 21, 2020 18:58
Show Gist options
  • Save dfaker/a8f0682d08481df0cb34dbe080090a9b to your computer and use it in GitHub Desktop.
Save dfaker/a8f0682d08481df0cb34dbe080090a9b to your computer and use it in GitHub Desktop.
Crop rect detection head for vgg19
import csv
csvname = 'frames\\coords.csv'
data = []
with open(csvname,'r', newline='') as csvfile:
reader = csv.DictReader(csvfile)
for row in reader:
data.append( row )
import os
root = 'frames'
frames = [os.path.join(root,x) for x in os.listdir(root)]
from keras.applications.vgg19 import VGG19
from keras.applications.vgg19 import preprocess_input
from keras.models import Model
from keras.layers import Dense, GlobalAveragePooling2D,Flatten
from keras.optimizers import SGD
import random
import cv2
import numpy as np
base_model = VGG19(weights='imagenet', include_top=False, input_shape=(224,224,3))
for l in base_model.layers:
l.trainable=True
x = base_model.output
x = GlobalAveragePooling2D()(x)
x = Dense(1024, activation='relu')(x)
predictions = Dense(4, activation='sigmoid')(x)
model = Model(inputs=base_model.input, outputs=predictions)
import tensorflow as tf
from keras import backend as K
class IoUloss:
def __call__(self,y_true, y_pred,sample_weight=0):
true_xy,true_wh = tf.split(y_true,[2,2],axis=-1)
pred_xy,pred_wh = tf.split(y_pred,[2,2],axis=-1)
true_wh_half = true_wh / 2.
true_mins = true_xy - true_wh_half
true_maxes = true_xy + true_wh_half
pred_wh_half = pred_wh / 2.
pred_mins = pred_xy - pred_wh_half
pred_maxes = pred_xy + pred_wh_half
intersect_mins = tf.maximum(pred_mins, true_mins)
intersect_maxes = tf.minimum(pred_maxes, true_maxes)
intersect_wh = intersect_maxes - intersect_mins
intersect_areas = intersect_wh[..., 0] * intersect_wh[..., 1]
true_areas = true_wh[..., 0] * true_wh[..., 1]
pred_areas = pred_wh[..., 0] * pred_wh[..., 1]
union_areas = pred_areas + true_areas - intersect_areas
iou_scores = tf.truediv(intersect_areas, union_areas)
return K.mean((1-iou_scores),axis=-1)
loss = IoUloss()
sgd = SGD(lr=1e-3, decay=1e-6, momentum=0.9, nesterov=True)
model.compile(optimizer=sgd, loss='mse')
try:
model.load_weights('cropweights.bin')
except Exception as e:
print(e)
imgCache = {}
from keras.preprocessing.image import ImageDataGenerator
data_gen_args = dict(
rescale=1./255,
shear_range=0.2,
rotation_range=15,
width_shift_range=0.1,
zoom_range=0.2,
height_shift_range=0.1,
brightness_range=(0.8,1.2),
channel_shift_range=0.5
)
image_datagen = ImageDataGenerator(**data_gen_args)
passn=0
samp=None
while 1:
passn+=1
xs=[]
y=[]
for sample in random.sample(data,min(len(data),30)):
imgsamp = imgCache.setdefault(sample['filename'],cv2.imread(sample['filename']))
seed = random.randint(0,10000)
imgsamp = image_datagen.random_transform(imgsamp,seed=seed)
xs.append( imgsamp )
y.append( [
float(sample['xc']),
float(sample['yc']),
float(sample['w' ]),
float(sample['h' ])
] )
x=preprocess_input(np.array(xs))
y=np.array(y)
loss = model.train_on_batch(x,y)
if passn%10==0 or samp is None:
samp = cv2.imread(random.choice(frames))
if passn%100==0:
model.save_weights('cropweights.bin')
pred = model.predict( preprocess_input(np.array([samp])) )[0]
xc,yc,w,h = pred
print(passn,loss,xc,yc,w,h)
oimg = samp.copy()
x0 = int((xc-(w/2))*224)
y0 = int((yc-(h/2))*224)
x1 = int((xc+(w/2))*224)
y1 = int((yc+(h/2))*224)
cv2.rectangle(oimg,(x0,y0),(x1,y1),(255,255,0),2 )
cv2.imshow('im',oimg)
cv2.waitKey(1)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment