diff --git a/.model_inference.py.swp b/.model_inference.py.swp index 99574fe8..3a921793 100644 Binary files a/.model_inference.py.swp and b/.model_inference.py.swp differ diff --git a/model_inference.py b/model_inference.py index 5999039b..1cdc7a6c 100644 --- a/model_inference.py +++ b/model_inference.py @@ -96,12 +96,12 @@ model_embedding.load_state_dict(torch.load('./embedding-output/model_embedding.p model_embedding.to(device) model_embedding.eval() -with open('/model_classifier.obj','rb') as file: +with open('./model_classifier.obj','rb') as file: model_classifier = pickle.load(file) classes = model_classifier.__getstate__()['classes_'] -video = Path('/content/test_videos_2022/2022-NLS-5-NLS_05_2022_Heli_UHD_01-000140-000155-Karussell.mp4') +video = Path('./test_videos_2022/2022-NLS-5-NLS_05_2022_Heli_UHD_01-000140-000155-Karussell.mp4') reader = cv2.VideoCapture(str(video)) fps = reader.get(cv2.CAP_PROP_FPS) @@ -109,18 +109,19 @@ w = int(reader.get(cv2.CAP_PROP_FRAME_WIDTH)) h = int(reader.get(cv2.CAP_PROP_FRAME_HEIGHT)) reader.release() -imgsz = check_img_size((w, h), s=model.stride) -dataset = LoadImages(video, img_size=imgsz, stride=model.stride, auto=model.pt) - weights_path = Path('./yolov5/best.pt') model = DetectMultiBackend(weights_path, device=torch.device(device)) +imgsz = check_img_size((w, h), s=model.stride) +dataset = LoadImages(video, img_size=imgsz, stride=model.stride, auto=model.pt) + save_dir = Path('./detection-output/') os.makedirs(save_dir) writer = cv2.VideoWriter(str(save_dir / 'res.mp4'), cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) for frame_n, (path, im, im0s, vid_cap, s) in enumerate(dataset): + print(frame_n) im = torch.from_numpy(im).to(device) im = im.half() if model.fp16 else im.float() # uint8 to fp16/32 im /= 255 # 0 - 255 to 0.0 - 1.0 diff --git a/requirements.txt b/requirements.txt index 7795139c..7eb91235 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ gdown +seaborn torch torchvision sklearn