Scarecrow / original_app /backend.py
Rgascoin's picture
Add hugging files
edcb323
raw
history blame
2.54 kB
import cv2
import numpy as np
import socket
import pickle
import struct
# Load YOLO model
net = cv2.dnn.readNet("yolov3.weights", "yolov3.cfg")
classes = []
with open("coco.names", "r") as f:
classes = [line.strip() for line in f.readlines()]
resolved_label = ''
# Set up socket
HOST = ''
PORT = 8089
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
print('Socket created')
s.bind((HOST, PORT))
print('Socket bind complete')
s.listen(10)
print('Socket now listening')
# Accept connections
conn, addr = s.accept()
# Receive and process frames
data = b''
payload_size = struct.calcsize("L")
while True:
# Retrieve message size
while len(data) < payload_size:
data += conn.recv(4096)
packed_msg_size = data[:payload_size]
data = data[payload_size:]
msg_size = struct.unpack("L", packed_msg_size)[0]
# Retrieve all data based on message size
while len(data) < msg_size:
data += conn.recv(4096)
frame_data = data[:msg_size]
data = data[msg_size:]
# Extract frame
frame = pickle.loads(frame_data)
# Run YOLO on frame
blob = cv2.dnn.blobFromImage(frame, 1/255.0, (416, 416), swapRB=True, crop=False)
net.setInput(blob)
outputs = net.forward(net.getUnconnectedOutLayersNames())
boxes = []
confidences = []
class_ids = []
for output in outputs:
for detection in output:
scores = detection[5:]
class_id = np.argmax(scores)
confidence = scores[class_id]
if confidence > 0.5:
center_x = int(detection[0] * frame.shape[1])
center_y = int(detection[1] * frame.shape[0])
w = int(detection[2] * frame.shape[1])
h = int(detection[3] * frame.shape[0])
x = int(center_x - w/2)
y = int(center_y - h/2)
boxes.append([x, y, w, h])
confidences.append(float(confidence))
class_ids.append(class_id)
indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)
if len(indexes) > 0:
for i in indexes.flatten():
resolved_label = classes[class_ids[i]]
print(resolved_label)
# Display frame
cv2.imshow('frame', frame)
cv2.waitKey(1)
# Send response to client
try:
if len(indexes) > 0:
response = "[Scarecrow]: " + resolved_label
else:
response = "[Scarecrow]: NONE"
except IndexError:
response = "[Scarecrow]: ERROR"
conn.sendall(response.encode())