|
|
@@ -182,15 +182,7 @@ class Displayer(QMainWindow):
|
|
|
|
|
|
def cv2_frame_to_cuda(frame, datatype = torch.float32):
|
|
|
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
|
- pic = Image.fromarray(frame)
|
|
|
- img = torch.ByteTensor(torch.ByteStorage.from_buffer(pic.tobytes()))
|
|
|
- img = img.cuda()
|
|
|
- img = img.view(pic.size[1], pic.size[0], len(pic.getbands()))
|
|
|
- img = img.permute((2, 0, 1)).contiguous()
|
|
|
- tmp = img.to(dtype=datatype).div(255)
|
|
|
- tmp.unsqueeze_(0)
|
|
|
- tmp = tmp.to(datatype)
|
|
|
- return tmp
|
|
|
+ return torch.as_tensor(frame, device=torch.device("cuda")).type(datatype).unsqueeze(0).permute(0, 3, 1, 2)/255
|
|
|
|
|
|
# --------------- Main ---------------
|
|
|
|
