File size: 1,694 Bytes
9e14838 | 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 | import cv2
class SupplyWriter:
def __init__(self, intput_video, output_video, opt_thres, rgb_input=True):
reader = cv2.VideoCapture(intput_video)
fourcc = cv2.VideoWriter_fourcc(*"XVID")
fps = reader.get(cv2.CAP_PROP_FPS)
width = int(reader.get(3))
height = int(reader.get(4))
reader.release()
self.padding = 40
self.writer = cv2.VideoWriter(output_video, fourcc, fps, (height, width)[::-1])
self.rgb_input = rgb_input
self.opt_thres = opt_thres
def run(self, images, scores, boxes):
# Text variables
font_face = cv2.FONT_HERSHEY_SIMPLEX
thickness = 5
font_scale = 3
for image, score, box in zip(images, scores, boxes):
if self.rgb_input:
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
if box is not None:
label = "fake" if score > self.opt_thres else "real"
x1, y1, x2, y2 = box
x = int(x1)
y = int(y1)
w = int(x2 - x1)
h = int(y2 - y1)
color = (
(255, 255, 0) if label == "real" else (0, 255, 255)
) # BGR 255 0
cv2.putText(
image,
label,
(x, y + h + 68),
font_face,
font_scale,
color,
thickness,
2,
)
# draw box over face
cv2.rectangle(image, (x, y), (x + w, y + h), color, 10)
self.writer.write(image)
self.writer.release()
|