forked from Jack-XHP/LabPicV2-MaskRCNN
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathRunPredictionOnVideo.py
133 lines (125 loc) · 8.07 KB
/
RunPredictionOnVideo.py
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
# Run trained net on video to generate prediction and write to another video
# ...............................Imports..................................................................
import cv2
import numpy as np
import torch
import os
import detection
from Utils.Visual import ChemDemo
import time
from vidgear.gears import WriteGear
# import scipy.misc as misc
############################################Input parameters###################################################################################
# -------------------------------------Input parameters-----------------------------------------------------------------------
device = torch.device('cuda') # Use GPU or CPU for prediction (GPU faster but demend nvidia GPU and CUDA installed else set UseGPU to False)
FreezeBatchNormStatistics = False # wether to freeze the batch statics on prediction setting this true or false might change the prediction mostly False work better
OutEnding = "" # Add This to file name
video_folder = "F:\LM113Pod2_20200225"
# -----------------------------------------Location of the pretrain model-----------------------------------------------------------------------------------
Trained_model_path = r"model_400.pth"
##################################Load net###########################################################################################
# ---------------------Create and Initiate net and create optimizer------------------------------------------------------------------------------------
print("predicting subclasses")
model = detection.__dict__['maskrcnn_resnet50_fpn'](num_classes=5, pretrained=False, num_sub_cls=17)
print("loading trained model")
checkpoint = torch.load(Trained_model_path, map_location='cpu')
model.load_state_dict(checkpoint['model'])
demo = ChemDemo(model, device=device, confidence_threshold=0.6)
output_params = {"-vcodec": "libx264", "-crf": 30, "-preset": "ultrafast"}
for workstation in os.listdir(video_folder):
path = os.path.join(video_folder, workstation)
if not os.path.isdir(path): continue
for day in os.listdir(os.path.join(video_folder, workstation)):
path = os.path.join(video_folder, workstation, day)
if not os.path.isdir(path): continue
for hour in os.listdir(os.path.join(video_folder, workstation, day)):
path = os.path.join(video_folder, workstation,day, hour)
if not os.path.isdir(path): continue
for minute in os.listdir(os.path.join(video_folder, workstation,day, hour)):
if len(minute) != 6: continue
# ---------------------OPEN video-----------------------------------------------------------------------------------------------------
t1 = time.time()
InputVideo = os.path.join(video_folder, workstation, day, hour, minute)
print(InputVideo)
OutVideoMain = InputVideo[:-4] + "_MainClasses.mp4" # Output video that contain vessel filled liquid and solid
OutVideoAll = InputVideo[:-4] + "_AllClasses.mp4" # Output video that contain subclasses that have more then 5% of the image
cap = cv2.VideoCapture(InputVideo)
MainCatsVideoWriter = None
AllCatsVideoWriter = None
# --------------------Create output video---------------------------------------------------------------------------------
# -----------------------Read Frame one by one-----------------------------------------------------------------------
# Read until video is completed
# iii=0
k = 0
while (cap.isOpened()):
# if iii>3: break
# Capture frame-by-frame
# ..................Read and resize image...............................................................................
ret, Im = cap.read()
k += 1
if k % 10 != 0: continue
if ret == False: break
# Display the resulting frame
h, w, d = Im.shape
r = np.max([h, w])
if r > 840: # Image larger then 840X840 are shrinked (this is not essential, but the net results might degrade when using to large images
fr = 840 / r
Im = cv2.resize(Im, (int(w * fr), int(h * fr)))
h, w, d = Im.shape
if not (type(Im) is np.ndarray): continue
Imgs = [torch.tensor(np.transpose(Im, (2, 0, 1))) / 255.0]
# ................................Make Prediction.............................................................................................................
with torch.autograd.no_grad():
preditions = demo.predict_on_image(Imgs)
# ------------------------------------Display main classes on the image----------------------------------------------------------------------------------
num_obj = len(preditions['labels'])
MainCatName = ['Liquid', 'Vessel', 'Solid']
AllCatName = ['Vessel', 'V Label', 'V Cork', 'V Part', 'Ignore', 'Liquid GENERAL',
'Liquid Suspension', 'Foam', 'Gel', 'Solid GENERAL', 'Granular', 'Powder',
'Solid Bulk', 'Vapor', 'Other', 'Filled']
results = {name: Im.copy() for name in MainCatName}
sub_res = {name: Im.copy() for name in AllCatName}
font = cv2.FONT_HERSHEY_SIMPLEX
for cls in MainCatName:
cv2.putText(results[cls], cls, (int(w / 3), int(h / 6)), font, 2, (0, 255, 0), 2, cv2.LINE_AA)
for cls in AllCatName:
cv2.putText(sub_res[cls], cls, (int(w / 3), int(h / 6)), font, 2, (0, 255, 0), 2, cv2.LINE_AA)
for i in range(num_obj):
mask = preditions['masks'][i].numpy()
cls = MainCatName[preditions['labels'][i] - 1]
res_img = results[cls]
sub_class = preditions['sub_cls'][i].nonzero().flatten() - 1
sub_imgs = [sub_res[AllCatName[idx]] for idx in sub_class]
imgs = [res_img] + sub_imgs
for img in imgs:
img[:, :, 1] = img[:, :, 1] * (1 - mask)
img[:, :, 0] = img[:, :, 0] + 255 * mask
my = 2
mx = 2
OutMain = np.zeros([h * my, w * mx, 3], np.uint8)
OutMain[0:h, 0:w] = results['Vessel']
OutMain[h: 2 * h, 0:w] = results['Liquid']
OutMain[0:h, w:2 * w] = results['Solid']
OutMain[h:2 * h, w:2 * w] = Im
if MainCatsVideoWriter is None:
h, w, d = OutMain.shape
MainCatsVideoWriter = WriteGear(output_filename=OutVideoMain, **output_params)
MainCatsVideoWriter.write(OutMain)
h, w, d = Im.shape
my = 3
mx = 3
OutMain = np.zeros([h * my, w * mx, 3], np.uint8)
for i in range(9):
cls = AllCatName[i + 5]
res = sub_res[cls]
OutMain[i % 3 * h: (i % 3 + 1) * h, i // 3 * w: (i // 3 + 1) * w] = res
if AllCatsVideoWriter is None:
h, w, d = OutMain.shape
AllCatsVideoWriter = WriteGear(output_filename=OutVideoAll, **output_params)
AllCatsVideoWriter.write(OutMain)
# -----------------------------------------------------------------------------------------------------------------------------
t2 = time.time()
print(f"Finished t:{t2 -t1}")
AllCatsVideoWriter.close()
MainCatsVideoWriter.close()
cap.release()