Spaces:
Running
Running
File size: 6,272 Bytes
5e2f8af 3048dcc 5e2f8af 3048dcc 5e2f8af 6cbade1 3048dcc 6cbade1 3048dcc 5e2f8af 3048dcc 5e2f8af 3048dcc 5e2f8af 3048dcc 5e2f8af 3048dcc 5e2f8af 3048dcc 5e2f8af 3048dcc 5e2f8af 3048dcc 5e2f8af 3048dcc 5e2f8af 3048dcc 5e2f8af 3048dcc 5e2f8af 3048dcc bd98f0f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 |
import base64
from uuid import uuid4
import gradio as gr
from fastcore.all import *
from fastai.vision.all import *
import numpy as np
import timm
def parent_labels(o):
"Label `item` with the parent folder name."
return Path(o).parent.name.split(",")
class LabelSmoothingBCEWithLogitsLossFlat(BCEWithLogitsLossFlat):
def __init__(self, eps:float=0.1, **kwargs):
self.eps = eps
super().__init__(thresh=0.1, **kwargs)
def __call__(self, inp, targ, **kwargs):
targ_smooth = targ.float() * (1. - self.eps) + 0.5 * self.eps
return super().__call__(inp, targ_smooth, **kwargs)
learn = load_learner('models.pkl')
# set a new loss function with a threshold of 0.4 to remove more false positives
learn.loss_func = BCEWithLogitsLossFlat(thresh=0.4)
def predict_tags(image, vtt, threshold=0.4):
vtt = base64.b64decode(vtt.replace("data:text/vtt;base64,", ""))
sprite = PILImage.create(image)
offsets = []
times = []
images = []
frames = []
for i, (left, top, right, bottom, time_seconds) in enumerate(getVTToffsets(vtt)):
frames.append(i)
times.append(time_seconds)
offsets.append((left, top, right, bottom))
cut_frame = sprite.crop((left, top, left + right, top + bottom))
images.append(PILImage.create(np.asarray(cut_frame)))
# create dataset
threshold = threshold or 0.4
learn.loss_func = BCEWithLogitsLossFlat(thresh=threshold)
test_dl = learn.dls.test_dl(images, bs=64)
# get predictions
probabilities, _, activations = learn.get_preds(dl=test_dl, with_decoded=True)
learn.loss_func = BCEWithLogitsLossFlat(thresh=0.4)
# swivel into tags list from activations
tags = {}
for idx1, activation in enumerate(activations):
for idx2, i in enumerate(activation):
if not i:
continue
tag = learn.dls.vocab[idx2]
tag = tag.replace("_", " ")
if tag not in tags:
tags[tag] = {'prob': 0, 'offset': (), 'frame': 0}
prob = float(probabilities[idx1][idx2])
if tags[tag]['prob'] < prob:
tags[tag]['prob'] = prob
tags[tag]['offset'] = offsets[idx1]
tags[tag]['frame'] = idx1
tags[tag]['time'] = times[idx1]
return tags
def predict_markers(image, vtt, threshold=0.4):
vtt = base64.b64decode(vtt.replace("data:text/vtt;base64,", ""))
sprite = PILImage.create(image)
offsets = []
times = []
images = []
frames = []
for i, (left, top, right, bottom, time_seconds) in enumerate(getVTToffsets(vtt)):
frames.append(i)
times.append(time_seconds)
offsets.append((left, top, right, bottom))
cut_frame = sprite.crop((left, top, left + right, top + bottom))
images.append(PILImage.create(np.asarray(cut_frame)))
# create dataset
threshold = threshold or 0.4
learn.loss_func = BCEWithLogitsLossFlat(thresh=threshold)
test_dl = learn.dls.test_dl(images, bs=64)
# get predictions
probabilities, _, activations = learn.get_preds(dl=test_dl, with_decoded=True)
learn.loss_func = BCEWithLogitsLossFlat(thresh=0.4)
# swivel into tags list from activations
all_data_per_frame = []
for idx1, activation in enumerate(activations):
frame_data = {'offset': offsets[idx1], 'frame': idx1, 'time': times[idx1], 'tags': []}
ftags = []
for idx2, i in enumerate(activation):
if not i:
continue
tag = learn.dls.vocab[idx2]
tag = tag.replace("_", " ")
prob = float(probabilities[idx1][idx2])
ftags.append({'label': tag, 'prob': prob})
if not ftags:
continue
frame_data['tags'] = ftags
all_data_per_frame.append(frame_data)
filtered = []
for idx, frame_data in enumerate(all_data_per_frame):
if idx == len(all_data_per_frame) - 1:
break
next_frame_data = all_data_per_frame[idx + 1]
frame_data['tags'] = [tag for tag in frame_data['tags'] for next_tag in next_frame_data['tags'] if tag['label'] == next_tag['label']]
if frame_data['tags']:
filtered.append(frame_data)
last_tag = set()
results = []
for frame_data in filtered:
tags = {s['label'] for s in frame_data['tags']}
if tags.intersection(last_tag):
continue
last_tag = tags
frame_data['tag'] = sorted(frame_data['tags'], key=lambda x: x['prob'], reverse=True)[0]
del frame_data['tags']
# add unique id to the frame
frame_data['id'] = str(uuid4())
results.append(frame_data)
return results
def getVTToffsets(vtt):
time_seconds = 0
left = top = right = bottom = None
for line in vtt.decode("utf-8").split("\n"):
line = line.strip()
if "-->" in line:
# grab the start time
# 00:00:00.000 --> 00:00:41.000
start = line.split("-->")[0].strip().split(":")
# convert to seconds
time_seconds = (
int(start[0]) * 3600
+ int(start[1]) * 60
+ float(start[2])
)
left = top = right = bottom = None
elif "xywh=" in line:
left, top, right, bottom = line.split("xywh=")[-1].split(",")
left, top, right, bottom = (
int(left),
int(top),
int(right),
int(bottom),
)
else:
continue
if not left:
continue
yield left, top, right, bottom, time_seconds
# create a gradio interface with 2 tabs
tag = gr.Interface(
fn=predict_tags,
inputs=[
gr.Image(),
gr.Textbox(label="VTT file"),
gr.Number(value=0.4, label="Threshold")
],
outputs=gr.JSON(label=""),
)
marker = gr.Interface(
fn=predict_markers,
inputs=[
gr.Image(),
gr.Textbox(label="VTT file"),
gr.Number(value=0.4, label="Threshold")
],
outputs=gr.JSON(label=""),
)
gr.TabbedInterface(
[tag, marker], ["tag", "marker"]
).launch(server_name="0.0.0.0")
|